stable-diffusion-webui

1
from collections import namedtuple
2
from copy import copy
3
from itertools import permutations, chain
4
import random
5
import csv
6
import os.path
7
from io import StringIO
8
from PIL import Image
9
import numpy as np
10

11
import modules.scripts as scripts
12
import gradio as gr
13

14
from modules import images, sd_samplers, processing, sd_models, sd_vae, sd_samplers_kdiffusion, errors
15
from modules.processing import process_images, Processed, StableDiffusionProcessingTxt2Img
16
from modules.shared import opts, state
17
import modules.shared as shared
18
import modules.sd_samplers
19
import modules.sd_models
20
import modules.sd_vae
21
import re
22

23
from modules.ui_components import ToolButton
24

25
fill_values_symbol = "\U0001f4d2"  # 📒
26

27
AxisInfo = namedtuple('AxisInfo', ['axis', 'values'])
28

29

30
def apply_field(field):
31
    def fun(p, x, xs):
32
        setattr(p, field, x)
33

34
    return fun
35

36

37
def apply_prompt(p, x, xs):
38
    if xs[0] not in p.prompt and xs[0] not in p.negative_prompt:
39
        raise RuntimeError(f"Prompt S/R did not find {xs[0]} in prompt or negative prompt.")
40

41
    p.prompt = p.prompt.replace(xs[0], x)
42
    p.negative_prompt = p.negative_prompt.replace(xs[0], x)
43

44

45
def apply_order(p, x, xs):
46
    token_order = []
47

48
    # Initally grab the tokens from the prompt, so they can be replaced in order of earliest seen
49
    for token in x:
50
        token_order.append((p.prompt.find(token), token))
51

52
    token_order.sort(key=lambda t: t[0])
53

54
    prompt_parts = []
55

56
    # Split the prompt up, taking out the tokens
57
    for _, token in token_order:
58
        n = p.prompt.find(token)
59
        prompt_parts.append(p.prompt[0:n])
60
        p.prompt = p.prompt[n + len(token):]
61

62
    # Rebuild the prompt with the tokens in the order we want
63
    prompt_tmp = ""
64
    for idx, part in enumerate(prompt_parts):
65
        prompt_tmp += part
66
        prompt_tmp += x[idx]
67
    p.prompt = prompt_tmp + p.prompt
68

69

70
def confirm_samplers(p, xs):
71
    for x in xs:
72
        if x.lower() not in sd_samplers.samplers_map:
73
            raise RuntimeError(f"Unknown sampler: {x}")
74

75

76
def apply_checkpoint(p, x, xs):
77
    info = modules.sd_models.get_closet_checkpoint_match(x)
78
    if info is None:
79
        raise RuntimeError(f"Unknown checkpoint: {x}")
80
    p.override_settings['sd_model_checkpoint'] = info.name
81

82

83
def confirm_checkpoints(p, xs):
84
    for x in xs:
85
        if modules.sd_models.get_closet_checkpoint_match(x) is None:
86
            raise RuntimeError(f"Unknown checkpoint: {x}")
87

88

89
def confirm_checkpoints_or_none(p, xs):
90
    for x in xs:
91
        if x in (None, "", "None", "none"):
92
            continue
93

94
        if modules.sd_models.get_closet_checkpoint_match(x) is None:
95
            raise RuntimeError(f"Unknown checkpoint: {x}")
96

97

98
def apply_clip_skip(p, x, xs):
99
    opts.data["CLIP_stop_at_last_layers"] = x
100

101

102
def apply_upscale_latent_space(p, x, xs):
103
    if x.lower().strip() != '0':
104
        opts.data["use_scale_latent_for_hires_fix"] = True
105
    else:
106
        opts.data["use_scale_latent_for_hires_fix"] = False
107

108

109
def find_vae(name: str):
110
    if name.lower() in ['auto', 'automatic']:
111
        return modules.sd_vae.unspecified
112
    if name.lower() == 'none':
113
        return None
114
    else:
115
        choices = [x for x in sorted(modules.sd_vae.vae_dict, key=lambda x: len(x)) if name.lower().strip() in x.lower()]
116
        if len(choices) == 0:
117
            print(f"No VAE found for {name}; using automatic")
118
            return modules.sd_vae.unspecified
119
        else:
120
            return modules.sd_vae.vae_dict[choices[0]]
121

122

123
def apply_vae(p, x, xs):
124
    modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file=find_vae(x))
125

126

127
def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _):
128
    p.styles.extend(x.split(','))
129

130

131
def apply_uni_pc_order(p, x, xs):
132
    opts.data["uni_pc_order"] = min(x, p.steps - 1)
133

134

135
def apply_face_restore(p, opt, x):
136
    opt = opt.lower()
137
    if opt == 'codeformer':
138
        is_active = True
139
        p.face_restoration_model = 'CodeFormer'
140
    elif opt == 'gfpgan':
141
        is_active = True
142
        p.face_restoration_model = 'GFPGAN'
143
    else:
144
        is_active = opt in ('true', 'yes', 'y', '1')
145

146
    p.restore_faces = is_active
147

148

149
def apply_override(field, boolean: bool = False):
150
    def fun(p, x, xs):
151
        if boolean:
152
            x = True if x.lower() == "true" else False
153
        p.override_settings[field] = x
154
    return fun
155

156

157
def boolean_choice(reverse: bool = False):
158
    def choice():
159
        return ["False", "True"] if reverse else ["True", "False"]
160
    return choice
161

162

163
def format_value_add_label(p, opt, x):
164
    if type(x) == float:
165
        x = round(x, 8)
166

167
    return f"{opt.label}: {x}"
168

169

170
def format_value(p, opt, x):
171
    if type(x) == float:
172
        x = round(x, 8)
173
    return x
174

175

176
def format_value_join_list(p, opt, x):
177
    return ", ".join(x)
178

179

180
def do_nothing(p, x, xs):
181
    pass
182

183

184
def format_nothing(p, opt, x):
185
    return ""
186

187

188
def format_remove_path(p, opt, x):
189
    return os.path.basename(x)
190

191

192
def str_permutations(x):
193
    """dummy function for specifying it in AxisOption's type when you want to get a list of permutations"""
194
    return x
195

196

197
def list_to_csv_string(data_list):
198
    with StringIO() as o:
199
        csv.writer(o).writerow(data_list)
200
        return o.getvalue().strip()
201

202

203
def csv_string_to_list_strip(data_str):
204
    return list(map(str.strip, chain.from_iterable(csv.reader(StringIO(data_str)))))
205

206

207
class AxisOption:
208
    def __init__(self, label, type, apply, format_value=format_value_add_label, confirm=None, cost=0.0, choices=None, prepare=None):
209
        self.label = label
210
        self.type = type
211
        self.apply = apply
212
        self.format_value = format_value
213
        self.confirm = confirm
214
        self.cost = cost
215
        self.prepare = prepare
216
        self.choices = choices
217

218

219
class AxisOptionImg2Img(AxisOption):
220
    def __init__(self, *args, **kwargs):
221
        super().__init__(*args, **kwargs)
222
        self.is_img2img = True
223

224

225
class AxisOptionTxt2Img(AxisOption):
226
    def __init__(self, *args, **kwargs):
227
        super().__init__(*args, **kwargs)
228
        self.is_img2img = False
229

230

231
axis_options = [
232
    AxisOption("Nothing", str, do_nothing, format_value=format_nothing),
233
    AxisOption("Seed", int, apply_field("seed")),
234
    AxisOption("Var. seed", int, apply_field("subseed")),
235
    AxisOption("Var. strength", float, apply_field("subseed_strength")),
236
    AxisOption("Steps", int, apply_field("steps")),
237
    AxisOptionTxt2Img("Hires steps", int, apply_field("hr_second_pass_steps")),
238
    AxisOption("CFG Scale", float, apply_field("cfg_scale")),
239
    AxisOptionImg2Img("Image CFG Scale", float, apply_field("image_cfg_scale")),
240
    AxisOption("Prompt S/R", str, apply_prompt, format_value=format_value),
241
    AxisOption("Prompt order", str_permutations, apply_order, format_value=format_value_join_list),
242
    AxisOptionTxt2Img("Sampler", str, apply_field("sampler_name"), format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers if x.name not in opts.hide_samplers]),
243
    AxisOptionTxt2Img("Hires sampler", str, apply_field("hr_sampler_name"), confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img if x.name not in opts.hide_samplers]),
244
    AxisOptionImg2Img("Sampler", str, apply_field("sampler_name"), format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img if x.name not in opts.hide_samplers]),
245
    AxisOption("Checkpoint name", str, apply_checkpoint, format_value=format_remove_path, confirm=confirm_checkpoints, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list, key=str.casefold)),
246
    AxisOption("Negative Guidance minimum sigma", float, apply_field("s_min_uncond")),
247
    AxisOption("Sigma Churn", float, apply_field("s_churn")),
248
    AxisOption("Sigma min", float, apply_field("s_tmin")),
249
    AxisOption("Sigma max", float, apply_field("s_tmax")),
250
    AxisOption("Sigma noise", float, apply_field("s_noise")),
251
    AxisOption("Schedule type", str, apply_override("k_sched_type"), choices=lambda: list(sd_samplers_kdiffusion.k_diffusion_scheduler)),
252
    AxisOption("Schedule min sigma", float, apply_override("sigma_min")),
253
    AxisOption("Schedule max sigma", float, apply_override("sigma_max")),
254
    AxisOption("Schedule rho", float, apply_override("rho")),
255
    AxisOption("Eta", float, apply_field("eta")),
256
    AxisOption("Clip skip", int, apply_clip_skip),
257
    AxisOption("Denoising", float, apply_field("denoising_strength")),
258
    AxisOption("Initial noise multiplier", float, apply_field("initial_noise_multiplier")),
259
    AxisOption("Extra noise", float, apply_override("img2img_extra_noise")),
260
    AxisOptionTxt2Img("Hires upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]),
261
    AxisOptionImg2Img("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight")),
262
    AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)),
263
    AxisOption("Styles", str, apply_styles, choices=lambda: list(shared.prompt_styles.styles)),
264
    AxisOption("UniPC Order", int, apply_uni_pc_order, cost=0.5),
265
    AxisOption("Face restore", str, apply_face_restore, format_value=format_value),
266
    AxisOption("Token merging ratio", float, apply_override('token_merging_ratio')),
267
    AxisOption("Token merging ratio high-res", float, apply_override('token_merging_ratio_hr')),
268
    AxisOption("Always discard next-to-last sigma", str, apply_override('always_discard_next_to_last_sigma', boolean=True), choices=boolean_choice(reverse=True)),
269
    AxisOption("SGM noise multiplier", str, apply_override('sgm_noise_multiplier', boolean=True), choices=boolean_choice(reverse=True)),
270
    AxisOption("Refiner checkpoint", str, apply_field('refiner_checkpoint'), format_value=format_remove_path, confirm=confirm_checkpoints_or_none, cost=1.0, choices=lambda: ['None'] + sorted(sd_models.checkpoints_list, key=str.casefold)),
271
    AxisOption("Refiner switch at", float, apply_field('refiner_switch_at')),
272
    AxisOption("RNG source", str, apply_override("randn_source"), choices=lambda: ["GPU", "CPU", "NV"]),
273
    AxisOption("FP8 mode", str, apply_override("fp8_storage"), cost=0.9, choices=lambda: ["Disable", "Enable for SDXL", "Enable"]),
274
]
275

276

277
def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size):
278
    hor_texts = [[images.GridAnnotation(x)] for x in x_labels]
279
    ver_texts = [[images.GridAnnotation(y)] for y in y_labels]
280
    title_texts = [[images.GridAnnotation(z)] for z in z_labels]
281

282
    list_size = (len(xs) * len(ys) * len(zs))
283

284
    processed_result = None
285

286
    state.job_count = list_size * p.n_iter
287

288
    def process_cell(x, y, z, ix, iy, iz):
289
        nonlocal processed_result
290

291
        def index(ix, iy, iz):
292
            return ix + iy * len(xs) + iz * len(xs) * len(ys)
293

294
        state.job = f"{index(ix, iy, iz) + 1} out of {list_size}"
295

296
        processed: Processed = cell(x, y, z, ix, iy, iz)
297

298
        if processed_result is None:
299
            # Use our first processed result object as a template container to hold our full results
300
            processed_result = copy(processed)
301
            processed_result.images = [None] * list_size
302
            processed_result.all_prompts = [None] * list_size
303
            processed_result.all_seeds = [None] * list_size
304
            processed_result.infotexts = [None] * list_size
305
            processed_result.index_of_first_image = 1
306

307
        idx = index(ix, iy, iz)
308
        if processed.images:
309
            # Non-empty list indicates some degree of success.
310
            processed_result.images[idx] = processed.images[0]
311
            processed_result.all_prompts[idx] = processed.prompt
312
            processed_result.all_seeds[idx] = processed.seed
313
            processed_result.infotexts[idx] = processed.infotexts[0]
314
        else:
315
            cell_mode = "P"
316
            cell_size = (processed_result.width, processed_result.height)
317
            if processed_result.images[0] is not None:
318
                cell_mode = processed_result.images[0].mode
319
                # This corrects size in case of batches:
320
                cell_size = processed_result.images[0].size
321
            processed_result.images[idx] = Image.new(cell_mode, cell_size)
322

323
    if first_axes_processed == 'x':
324
        for ix, x in enumerate(xs):
325
            if second_axes_processed == 'y':
326
                for iy, y in enumerate(ys):
327
                    for iz, z in enumerate(zs):
328
                        process_cell(x, y, z, ix, iy, iz)
329
            else:
330
                for iz, z in enumerate(zs):
331
                    for iy, y in enumerate(ys):
332
                        process_cell(x, y, z, ix, iy, iz)
333
    elif first_axes_processed == 'y':
334
        for iy, y in enumerate(ys):
335
            if second_axes_processed == 'x':
336
                for ix, x in enumerate(xs):
337
                    for iz, z in enumerate(zs):
338
                        process_cell(x, y, z, ix, iy, iz)
339
            else:
340
                for iz, z in enumerate(zs):
341
                    for ix, x in enumerate(xs):
342
                        process_cell(x, y, z, ix, iy, iz)
343
    elif first_axes_processed == 'z':
344
        for iz, z in enumerate(zs):
345
            if second_axes_processed == 'x':
346
                for ix, x in enumerate(xs):
347
                    for iy, y in enumerate(ys):
348
                        process_cell(x, y, z, ix, iy, iz)
349
            else:
350
                for iy, y in enumerate(ys):
351
                    for ix, x in enumerate(xs):
352
                        process_cell(x, y, z, ix, iy, iz)
353

354
    if not processed_result:
355
        # Should never happen, I've only seen it on one of four open tabs and it needed to refresh.
356
        print("Unexpected error: Processing could not begin, you may need to refresh the tab or restart the service.")
357
        return Processed(p, [])
358
    elif not any(processed_result.images):
359
        print("Unexpected error: draw_xyz_grid failed to return even a single processed image")
360
        return Processed(p, [])
361

362
    z_count = len(zs)
363

364
    for i in range(z_count):
365
        start_index = (i * len(xs) * len(ys)) + i
366
        end_index = start_index + len(xs) * len(ys)
367
        grid = images.image_grid(processed_result.images[start_index:end_index], rows=len(ys))
368
        if draw_legend:
369
            grid = images.draw_grid_annotations(grid, processed_result.images[start_index].size[0], processed_result.images[start_index].size[1], hor_texts, ver_texts, margin_size)
370
        processed_result.images.insert(i, grid)
371
        processed_result.all_prompts.insert(i, processed_result.all_prompts[start_index])
372
        processed_result.all_seeds.insert(i, processed_result.all_seeds[start_index])
373
        processed_result.infotexts.insert(i, processed_result.infotexts[start_index])
374

375
    sub_grid_size = processed_result.images[0].size
376
    z_grid = images.image_grid(processed_result.images[:z_count], rows=1)
377
    if draw_legend:
378
        z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], title_texts, [[images.GridAnnotation()]])
379
    processed_result.images.insert(0, z_grid)
380
    # TODO: Deeper aspects of the program rely on grid info being misaligned between metadata arrays, which is not ideal.
381
    # processed_result.all_prompts.insert(0, processed_result.all_prompts[0])
382
    # processed_result.all_seeds.insert(0, processed_result.all_seeds[0])
383
    processed_result.infotexts.insert(0, processed_result.infotexts[0])
384

385
    return processed_result
386

387

388
class SharedSettingsStackHelper(object):
389
    def __enter__(self):
390
        self.CLIP_stop_at_last_layers = opts.CLIP_stop_at_last_layers
391
        self.vae = opts.sd_vae
392
        self.uni_pc_order = opts.uni_pc_order
393

394
    def __exit__(self, exc_type, exc_value, tb):
395
        opts.data["sd_vae"] = self.vae
396
        opts.data["uni_pc_order"] = self.uni_pc_order
397
        modules.sd_models.reload_model_weights()
398
        modules.sd_vae.reload_vae_weights()
399

400
        opts.data["CLIP_stop_at_last_layers"] = self.CLIP_stop_at_last_layers
401

402

403
re_range = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\(([+-]\d+)\s*\))?\s*")
404
re_range_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+(?:.\d*)?)(?:\s*\(([+-]\d+(?:.\d*)?)\s*\))?\s*")
405

406
re_range_count = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\[(\d+)\s*])?\s*")
407
re_range_count_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+(?:.\d*)?)(?:\s*\[(\d+(?:.\d*)?)\s*])?\s*")
408

409

410
class Script(scripts.Script):
411
    def title(self):
412
        return "X/Y/Z plot"
413

414
    def ui(self, is_img2img):
415
        self.current_axis_options = [x for x in axis_options if type(x) == AxisOption or x.is_img2img == is_img2img]
416

417
        with gr.Row():
418
            with gr.Column(scale=19):
419
                with gr.Row():
420
                    x_type = gr.Dropdown(label="X type", choices=[x.label for x in self.current_axis_options], value=self.current_axis_options[1].label, type="index", elem_id=self.elem_id("x_type"))
421
                    x_values = gr.Textbox(label="X values", lines=1, elem_id=self.elem_id("x_values"))
422
                    x_values_dropdown = gr.Dropdown(label="X values", visible=False, multiselect=True, interactive=True)
423
                    fill_x_button = ToolButton(value=fill_values_symbol, elem_id="xyz_grid_fill_x_tool_button", visible=False)
424

425
                with gr.Row():
426
                    y_type = gr.Dropdown(label="Y type", choices=[x.label for x in self.current_axis_options], value=self.current_axis_options[0].label, type="index", elem_id=self.elem_id("y_type"))
427
                    y_values = gr.Textbox(label="Y values", lines=1, elem_id=self.elem_id("y_values"))
428
                    y_values_dropdown = gr.Dropdown(label="Y values", visible=False, multiselect=True, interactive=True)
429
                    fill_y_button = ToolButton(value=fill_values_symbol, elem_id="xyz_grid_fill_y_tool_button", visible=False)
430

431
                with gr.Row():
432
                    z_type = gr.Dropdown(label="Z type", choices=[x.label for x in self.current_axis_options], value=self.current_axis_options[0].label, type="index", elem_id=self.elem_id("z_type"))
433
                    z_values = gr.Textbox(label="Z values", lines=1, elem_id=self.elem_id("z_values"))
434
                    z_values_dropdown = gr.Dropdown(label="Z values", visible=False, multiselect=True, interactive=True)
435
                    fill_z_button = ToolButton(value=fill_values_symbol, elem_id="xyz_grid_fill_z_tool_button", visible=False)
436

437
        with gr.Row(variant="compact", elem_id="axis_options"):
438
            with gr.Column():
439
                draw_legend = gr.Checkbox(label='Draw legend', value=True, elem_id=self.elem_id("draw_legend"))
440
                no_fixed_seeds = gr.Checkbox(label='Keep -1 for seeds', value=False, elem_id=self.elem_id("no_fixed_seeds"))
441
                with gr.Row():
442
                    vary_seeds_x = gr.Checkbox(label='Vary seeds for X', value=False, min_width=80, elem_id=self.elem_id("vary_seeds_x"), tooltip="Use different seeds for images along X axis.")
443
                    vary_seeds_y = gr.Checkbox(label='Vary seeds for Y', value=False, min_width=80, elem_id=self.elem_id("vary_seeds_y"), tooltip="Use different seeds for images along Y axis.")
444
                    vary_seeds_z = gr.Checkbox(label='Vary seeds for Z', value=False, min_width=80, elem_id=self.elem_id("vary_seeds_z"), tooltip="Use different seeds for images along Z axis.")
445
            with gr.Column():
446
                include_lone_images = gr.Checkbox(label='Include Sub Images', value=False, elem_id=self.elem_id("include_lone_images"))
447
                include_sub_grids = gr.Checkbox(label='Include Sub Grids', value=False, elem_id=self.elem_id("include_sub_grids"))
448
                csv_mode = gr.Checkbox(label='Use text inputs instead of dropdowns', value=False, elem_id=self.elem_id("csv_mode"))
449
            with gr.Column():
450
                margin_size = gr.Slider(label="Grid margins (px)", minimum=0, maximum=500, value=0, step=2, elem_id=self.elem_id("margin_size"))
451

452
        with gr.Row(variant="compact", elem_id="swap_axes"):
453
            swap_xy_axes_button = gr.Button(value="Swap X/Y axes", elem_id="xy_grid_swap_axes_button")
454
            swap_yz_axes_button = gr.Button(value="Swap Y/Z axes", elem_id="yz_grid_swap_axes_button")
455
            swap_xz_axes_button = gr.Button(value="Swap X/Z axes", elem_id="xz_grid_swap_axes_button")
456

457
        def swap_axes(axis1_type, axis1_values, axis1_values_dropdown, axis2_type, axis2_values, axis2_values_dropdown):
458
            return self.current_axis_options[axis2_type].label, axis2_values, axis2_values_dropdown, self.current_axis_options[axis1_type].label, axis1_values, axis1_values_dropdown
459

460
        xy_swap_args = [x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown]
461
        swap_xy_axes_button.click(swap_axes, inputs=xy_swap_args, outputs=xy_swap_args)
462
        yz_swap_args = [y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown]
463
        swap_yz_axes_button.click(swap_axes, inputs=yz_swap_args, outputs=yz_swap_args)
464
        xz_swap_args = [x_type, x_values, x_values_dropdown, z_type, z_values, z_values_dropdown]
465
        swap_xz_axes_button.click(swap_axes, inputs=xz_swap_args, outputs=xz_swap_args)
466

467
        def fill(axis_type, csv_mode):
468
            axis = self.current_axis_options[axis_type]
469
            if axis.choices:
470
                if csv_mode:
471
                    return list_to_csv_string(axis.choices()), gr.update()
472
                else:
473
                    return gr.update(), axis.choices()
474
            else:
475
                return gr.update(), gr.update()
476

477
        fill_x_button.click(fn=fill, inputs=[x_type, csv_mode], outputs=[x_values, x_values_dropdown])
478
        fill_y_button.click(fn=fill, inputs=[y_type, csv_mode], outputs=[y_values, y_values_dropdown])
479
        fill_z_button.click(fn=fill, inputs=[z_type, csv_mode], outputs=[z_values, z_values_dropdown])
480

481
        def select_axis(axis_type, axis_values, axis_values_dropdown, csv_mode):
482
            axis_type = axis_type or 0  # if axle type is None set to 0
483

484
            choices = self.current_axis_options[axis_type].choices
485
            has_choices = choices is not None
486

487
            if has_choices:
488
                choices = choices()
489
                if csv_mode:
490
                    if axis_values_dropdown:
491
                        axis_values = list_to_csv_string(list(filter(lambda x: x in choices, axis_values_dropdown)))
492
                        axis_values_dropdown = []
493
                else:
494
                    if axis_values:
495
                        axis_values_dropdown = list(filter(lambda x: x in choices, csv_string_to_list_strip(axis_values)))
496
                        axis_values = ""
497

498
            return (gr.Button.update(visible=has_choices), gr.Textbox.update(visible=not has_choices or csv_mode, value=axis_values),
499
                    gr.update(choices=choices if has_choices else None, visible=has_choices and not csv_mode, value=axis_values_dropdown))
500

501
        x_type.change(fn=select_axis, inputs=[x_type, x_values, x_values_dropdown, csv_mode], outputs=[fill_x_button, x_values, x_values_dropdown])
502
        y_type.change(fn=select_axis, inputs=[y_type, y_values, y_values_dropdown, csv_mode], outputs=[fill_y_button, y_values, y_values_dropdown])
503
        z_type.change(fn=select_axis, inputs=[z_type, z_values, z_values_dropdown, csv_mode], outputs=[fill_z_button, z_values, z_values_dropdown])
504

505
        def change_choice_mode(csv_mode, x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown):
506
            _fill_x_button, _x_values, _x_values_dropdown = select_axis(x_type, x_values, x_values_dropdown, csv_mode)
507
            _fill_y_button, _y_values, _y_values_dropdown = select_axis(y_type, y_values, y_values_dropdown, csv_mode)
508
            _fill_z_button, _z_values, _z_values_dropdown = select_axis(z_type, z_values, z_values_dropdown, csv_mode)
509
            return _fill_x_button, _x_values, _x_values_dropdown, _fill_y_button, _y_values, _y_values_dropdown, _fill_z_button, _z_values, _z_values_dropdown
510

511
        csv_mode.change(fn=change_choice_mode, inputs=[csv_mode, x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown], outputs=[fill_x_button, x_values, x_values_dropdown, fill_y_button, y_values, y_values_dropdown, fill_z_button, z_values, z_values_dropdown])
512

513
        def get_dropdown_update_from_params(axis, params):
514
            val_key = f"{axis} Values"
515
            vals = params.get(val_key, "")
516
            valslist = csv_string_to_list_strip(vals)
517
            return gr.update(value=valslist)
518

519
        self.infotext_fields = (
520
            (x_type, "X Type"),
521
            (x_values, "X Values"),
522
            (x_values_dropdown, lambda params: get_dropdown_update_from_params("X", params)),
523
            (y_type, "Y Type"),
524
            (y_values, "Y Values"),
525
            (y_values_dropdown, lambda params: get_dropdown_update_from_params("Y", params)),
526
            (z_type, "Z Type"),
527
            (z_values, "Z Values"),
528
            (z_values_dropdown, lambda params: get_dropdown_update_from_params("Z", params)),
529
        )
530

531
        return [x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown, draw_legend, include_lone_images, include_sub_grids, no_fixed_seeds, vary_seeds_x, vary_seeds_y, vary_seeds_z, margin_size, csv_mode]
532

533
    def run(self, p, x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown, draw_legend, include_lone_images, include_sub_grids, no_fixed_seeds, vary_seeds_x, vary_seeds_y, vary_seeds_z, margin_size, csv_mode):
534
        x_type, y_type, z_type = x_type or 0, y_type or 0, z_type or 0  # if axle type is None set to 0
535

536
        if not no_fixed_seeds:
537
            modules.processing.fix_seed(p)
538

539
        if not opts.return_grid:
540
            p.batch_size = 1
541

542
        def process_axis(opt, vals, vals_dropdown):
543
            if opt.label == 'Nothing':
544
                return [0]
545

546
            if opt.choices is not None and not csv_mode:
547
                valslist = vals_dropdown
548
            elif opt.prepare is not None:
549
                valslist = opt.prepare(vals)
550
            else:
551
                valslist = csv_string_to_list_strip(vals)
552

553
            if opt.type == int:
554
                valslist_ext = []
555

556
                for val in valslist:
557
                    if val.strip() == '':
558
                        continue
559
                    m = re_range.fullmatch(val)
560
                    mc = re_range_count.fullmatch(val)
561
                    if m is not None:
562
                        start = int(m.group(1))
563
                        end = int(m.group(2))+1
564
                        step = int(m.group(3)) if m.group(3) is not None else 1
565

566
                        valslist_ext += list(range(start, end, step))
567
                    elif mc is not None:
568
                        start = int(mc.group(1))
569
                        end = int(mc.group(2))
570
                        num = int(mc.group(3)) if mc.group(3) is not None else 1
571

572
                        valslist_ext += [int(x) for x in np.linspace(start=start, stop=end, num=num).tolist()]
573
                    else:
574
                        valslist_ext.append(val)
575

576
                valslist = valslist_ext
577
            elif opt.type == float:
578
                valslist_ext = []
579

580
                for val in valslist:
581
                    if val.strip() == '':
582
                        continue
583
                    m = re_range_float.fullmatch(val)
584
                    mc = re_range_count_float.fullmatch(val)
585
                    if m is not None:
586
                        start = float(m.group(1))
587
                        end = float(m.group(2))
588
                        step = float(m.group(3)) if m.group(3) is not None else 1
589

590
                        valslist_ext += np.arange(start, end + step, step).tolist()
591
                    elif mc is not None:
592
                        start = float(mc.group(1))
593
                        end = float(mc.group(2))
594
                        num = int(mc.group(3)) if mc.group(3) is not None else 1
595

596
                        valslist_ext += np.linspace(start=start, stop=end, num=num).tolist()
597
                    else:
598
                        valslist_ext.append(val)
599

600
                valslist = valslist_ext
601
            elif opt.type == str_permutations:
602
                valslist = list(permutations(valslist))
603

604
            valslist = [opt.type(x) for x in valslist]
605

606
            # Confirm options are valid before starting
607
            if opt.confirm:
608
                opt.confirm(p, valslist)
609

610
            return valslist
611

612
        x_opt = self.current_axis_options[x_type]
613
        if x_opt.choices is not None and not csv_mode:
614
            x_values = list_to_csv_string(x_values_dropdown)
615
        xs = process_axis(x_opt, x_values, x_values_dropdown)
616

617
        y_opt = self.current_axis_options[y_type]
618
        if y_opt.choices is not None and not csv_mode:
619
            y_values = list_to_csv_string(y_values_dropdown)
620
        ys = process_axis(y_opt, y_values, y_values_dropdown)
621

622
        z_opt = self.current_axis_options[z_type]
623
        if z_opt.choices is not None and not csv_mode:
624
            z_values = list_to_csv_string(z_values_dropdown)
625
        zs = process_axis(z_opt, z_values, z_values_dropdown)
626

627
        # this could be moved to common code, but unlikely to be ever triggered anywhere else
628
        Image.MAX_IMAGE_PIXELS = None  # disable check in Pillow and rely on check below to allow large custom image sizes
629
        grid_mp = round(len(xs) * len(ys) * len(zs) * p.width * p.height / 1000000)
630
        assert grid_mp < opts.img_max_size_mp, f'Error: Resulting grid would be too large ({grid_mp} MPixels) (max configured size is {opts.img_max_size_mp} MPixels)'
631

632
        def fix_axis_seeds(axis_opt, axis_list):
633
            if axis_opt.label in ['Seed', 'Var. seed']:
634
                return [int(random.randrange(4294967294)) if val is None or val == '' or val == -1 else val for val in axis_list]
635
            else:
636
                return axis_list
637

638
        if not no_fixed_seeds:
639
            xs = fix_axis_seeds(x_opt, xs)
640
            ys = fix_axis_seeds(y_opt, ys)
641
            zs = fix_axis_seeds(z_opt, zs)
642

643
        if x_opt.label == 'Steps':
644
            total_steps = sum(xs) * len(ys) * len(zs)
645
        elif y_opt.label == 'Steps':
646
            total_steps = sum(ys) * len(xs) * len(zs)
647
        elif z_opt.label == 'Steps':
648
            total_steps = sum(zs) * len(xs) * len(ys)
649
        else:
650
            total_steps = p.steps * len(xs) * len(ys) * len(zs)
651

652
        if isinstance(p, StableDiffusionProcessingTxt2Img) and p.enable_hr:
653
            if x_opt.label == "Hires steps":
654
                total_steps += sum(xs) * len(ys) * len(zs)
655
            elif y_opt.label == "Hires steps":
656
                total_steps += sum(ys) * len(xs) * len(zs)
657
            elif z_opt.label == "Hires steps":
658
                total_steps += sum(zs) * len(xs) * len(ys)
659
            elif p.hr_second_pass_steps:
660
                total_steps += p.hr_second_pass_steps * len(xs) * len(ys) * len(zs)
661
            else:
662
                total_steps *= 2
663

664
        total_steps *= p.n_iter
665

666
        image_cell_count = p.n_iter * p.batch_size
667
        cell_console_text = f"; {image_cell_count} images per cell" if image_cell_count > 1 else ""
668
        plural_s = 's' if len(zs) > 1 else ''
669
        print(f"X/Y/Z plot will create {len(xs) * len(ys) * len(zs) * image_cell_count} images on {len(zs)} {len(xs)}x{len(ys)} grid{plural_s}{cell_console_text}. (Total steps to process: {total_steps})")
670
        shared.total_tqdm.updateTotal(total_steps)
671

672
        state.xyz_plot_x = AxisInfo(x_opt, xs)
673
        state.xyz_plot_y = AxisInfo(y_opt, ys)
674
        state.xyz_plot_z = AxisInfo(z_opt, zs)
675

676
        # If one of the axes is very slow to change between (like SD model
677
        # checkpoint), then make sure it is in the outer iteration of the nested
678
        # `for` loop.
679
        first_axes_processed = 'z'
680
        second_axes_processed = 'y'
681
        if x_opt.cost > y_opt.cost and x_opt.cost > z_opt.cost:
682
            first_axes_processed = 'x'
683
            if y_opt.cost > z_opt.cost:
684
                second_axes_processed = 'y'
685
            else:
686
                second_axes_processed = 'z'
687
        elif y_opt.cost > x_opt.cost and y_opt.cost > z_opt.cost:
688
            first_axes_processed = 'y'
689
            if x_opt.cost > z_opt.cost:
690
                second_axes_processed = 'x'
691
            else:
692
                second_axes_processed = 'z'
693
        elif z_opt.cost > x_opt.cost and z_opt.cost > y_opt.cost:
694
            first_axes_processed = 'z'
695
            if x_opt.cost > y_opt.cost:
696
                second_axes_processed = 'x'
697
            else:
698
                second_axes_processed = 'y'
699

700
        grid_infotext = [None] * (1 + len(zs))
701

702
        def cell(x, y, z, ix, iy, iz):
703
            if shared.state.interrupted or state.stopping_generation:
704
                return Processed(p, [], p.seed, "")
705

706
            pc = copy(p)
707
            pc.styles = pc.styles[:]
708
            x_opt.apply(pc, x, xs)
709
            y_opt.apply(pc, y, ys)
710
            z_opt.apply(pc, z, zs)
711

712
            xdim = len(xs) if vary_seeds_x else 1
713
            ydim = len(ys) if vary_seeds_y else 1
714

715
            if vary_seeds_x:
716
               pc.seed += ix
717
            if vary_seeds_y:
718
               pc.seed += iy * xdim
719
            if vary_seeds_z:
720
               pc.seed += iz * xdim * ydim
721

722
            try:
723
                res = process_images(pc)
724
            except Exception as e:
725
                errors.display(e, "generating image for xyz plot")
726

727
                res = Processed(p, [], p.seed, "")
728

729
            # Sets subgrid infotexts
730
            subgrid_index = 1 + iz
731
            if grid_infotext[subgrid_index] is None and ix == 0 and iy == 0:
732
                pc.extra_generation_params = copy(pc.extra_generation_params)
733
                pc.extra_generation_params['Script'] = self.title()
734

735
                if x_opt.label != 'Nothing':
736
                    pc.extra_generation_params["X Type"] = x_opt.label
737
                    pc.extra_generation_params["X Values"] = x_values
738
                    if x_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:
739
                        pc.extra_generation_params["Fixed X Values"] = ", ".join([str(x) for x in xs])
740

741
                if y_opt.label != 'Nothing':
742
                    pc.extra_generation_params["Y Type"] = y_opt.label
743
                    pc.extra_generation_params["Y Values"] = y_values
744
                    if y_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:
745
                        pc.extra_generation_params["Fixed Y Values"] = ", ".join([str(y) for y in ys])
746

747
                grid_infotext[subgrid_index] = processing.create_infotext(pc, pc.all_prompts, pc.all_seeds, pc.all_subseeds)
748

749
            # Sets main grid infotext
750
            if grid_infotext[0] is None and ix == 0 and iy == 0 and iz == 0:
751
                pc.extra_generation_params = copy(pc.extra_generation_params)
752

753
                if z_opt.label != 'Nothing':
754
                    pc.extra_generation_params["Z Type"] = z_opt.label
755
                    pc.extra_generation_params["Z Values"] = z_values
756
                    if z_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:
757
                        pc.extra_generation_params["Fixed Z Values"] = ", ".join([str(z) for z in zs])
758

759
                grid_infotext[0] = processing.create_infotext(pc, pc.all_prompts, pc.all_seeds, pc.all_subseeds)
760

761
            return res
762

763
        with SharedSettingsStackHelper():
764
            processed = draw_xyz_grid(
765
                p,
766
                xs=xs,
767
                ys=ys,
768
                zs=zs,
769
                x_labels=[x_opt.format_value(p, x_opt, x) for x in xs],
770
                y_labels=[y_opt.format_value(p, y_opt, y) for y in ys],
771
                z_labels=[z_opt.format_value(p, z_opt, z) for z in zs],
772
                cell=cell,
773
                draw_legend=draw_legend,
774
                include_lone_images=include_lone_images,
775
                include_sub_grids=include_sub_grids,
776
                first_axes_processed=first_axes_processed,
777
                second_axes_processed=second_axes_processed,
778
                margin_size=margin_size
779
            )
780

781
        if not processed.images:
782
            # It broke, no further handling needed.
783
            return processed
784

785
        z_count = len(zs)
786

787
        # Set the grid infotexts to the real ones with extra_generation_params (1 main grid + z_count sub-grids)
788
        processed.infotexts[:1+z_count] = grid_infotext[:1+z_count]
789

790
        if not include_lone_images:
791
            # Don't need sub-images anymore, drop from list:
792
            processed.images = processed.images[:z_count+1]
793

794
        if opts.grid_save:
795
            # Auto-save main and sub-grids:
796
            grid_count = z_count + 1 if z_count > 1 else 1
797
            for g in range(grid_count):
798
                # TODO: See previous comment about intentional data misalignment.
799
                adj_g = g-1 if g > 0 else g
800
                images.save_image(processed.images[g], p.outpath_grids, "xyz_grid", info=processed.infotexts[g], extension=opts.grid_format, prompt=processed.all_prompts[adj_g], seed=processed.all_seeds[adj_g], grid=True, p=processed)
801
                if not include_sub_grids:  # if not include_sub_grids then skip saving after the first grid
802
                    break
803

804
        if not include_sub_grids:
805
            # Done with sub-grids, drop all related information:
806
            for _ in range(z_count):
807
                del processed.images[1]
808
                del processed.all_prompts[1]
809
                del processed.all_seeds[1]
810
                del processed.infotexts[1]
811

812
        return processed
813