1 files changed, 143 insertions, 82 deletions
diff --git a/scripts/xy_grid.py b/scripts/xy_grid.py
index f04d9b7e..6629f5d5 100644
--- a/scripts/xy_grid.py
+++ b/scripts/xy_grid.py
@@ -10,7 +10,7 @@ import numpy as np
 import modules.scripts as scripts
 import gradio as gr
 
-from modules import images, paths, sd_samplers, processing
+from modules import images, paths, sd_samplers, processing, sd_models, sd_vae
 from modules.hypernetworks import hypernetwork
 from modules.processing import process_images, Processed, StableDiffusionProcessingTxt2Img
 from modules.shared import opts, cmd_opts, state
@@ -22,6 +22,10 @@ import glob
 import os
 import re
 
+from modules.ui_components import ToolButton
+
+fill_values_symbol = "\U0001f4d2"  # 📒
+
 
 def apply_field(field):
     def fun(p, x, xs):
@@ -82,7 +86,6 @@ def apply_checkpoint(p, x, xs):
     if info is None:
         raise RuntimeError(f"Unknown checkpoint: {x}")
     modules.sd_models.reload_model_weights(shared.sd_model, info)
-    p.sd_model = shared.sd_model
 
 
 def confirm_checkpoints(p, xs):
@@ -125,24 +128,21 @@ def apply_upscale_latent_space(p, x, xs):
 
 
 def find_vae(name: str):
-    if name.lower() in ['auto', 'none']:
-        return name
+    if name.lower() in ['auto', 'automatic']:
+        return modules.sd_vae.unspecified
+    if name.lower() == 'none':
+        return None
     else:
-        vae_path = os.path.abspath(os.path.join(paths.models_path, 'VAE'))
-        found = glob.glob(os.path.join(vae_path, f'**/{name}.*pt'), recursive=True)
-        if found:
-            return found[0]
+        choices = [x for x in sorted(modules.sd_vae.vae_dict, key=lambda x: len(x)) if name.lower().strip() in x.lower()]
+        if len(choices) == 0:
+            print(f"No VAE found for {name}; using automatic")
+            return modules.sd_vae.unspecified
         else:
-            return 'auto'
+            return modules.sd_vae.vae_dict[choices[0]]
 
 
 def apply_vae(p, x, xs):
-    if x.lower().strip() == 'none':
-        modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file='None')
-    else:
-        found = find_vae(x)
-        if found:
-            v = modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file=found)
+    modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file=find_vae(x))
 
 
 def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _):
@@ -178,80 +178,106 @@ def str_permutations(x):
     """dummy function for specifying it in AxisOption's type when you want to get a list of permutations"""
     return x
 
-AxisOption = namedtuple("AxisOption", ["label", "type", "apply", "format_value", "confirm"])
-AxisOptionImg2Img = namedtuple("AxisOptionImg2Img", ["label", "type", "apply", "format_value", "confirm"])
+
+class AxisOption:
+    def __init__(self, label, type, apply, format_value=format_value_add_label, confirm=None, cost=0.0, choices=None):
+        self.label = label
+        self.type = type
+        self.apply = apply
+        self.format_value = format_value
+        self.confirm = confirm
+        self.cost = cost
+        self.choices = choices
+        self.is_img2img = False
+
+
+class AxisOptionImg2Img(AxisOption):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.is_img2img = False
 
 
 axis_options = [
-    AxisOption("Nothing", str, do_nothing, format_nothing, None),
-    AxisOption("Seed", int, apply_field("seed"), format_value_add_label, None),
-    AxisOption("Var. seed", int, apply_field("subseed"), format_value_add_label, None),
-    AxisOption("Var. strength", float, apply_field("subseed_strength"), format_value_add_label, None),
-    AxisOption("Steps", int, apply_field("steps"), format_value_add_label, None),
-    AxisOption("CFG Scale", float, apply_field("cfg_scale"), format_value_add_label, None),
-    AxisOption("Prompt S/R", str, apply_prompt, format_value, None),
-    AxisOption("Prompt order", str_permutations, apply_order, format_value_join_list, None),
-    AxisOption("Sampler", str, apply_sampler, format_value, confirm_samplers),
-    AxisOption("Checkpoint name", str, apply_checkpoint, format_value, confirm_checkpoints),
-    AxisOption("Hypernetwork", str, apply_hypernetwork, format_value, confirm_hypernetworks),
-    AxisOption("Hypernet str.", float, apply_hypernetwork_strength, format_value_add_label, None),
-    AxisOption("Sigma Churn", float, apply_field("s_churn"), format_value_add_label, None),
-    AxisOption("Sigma min", float, apply_field("s_tmin"), format_value_add_label, None),
-    AxisOption("Sigma max", float, apply_field("s_tmax"), format_value_add_label, None),
-    AxisOption("Sigma noise", float, apply_field("s_noise"), format_value_add_label, None),
-    AxisOption("Eta", float, apply_field("eta"), format_value_add_label, None),
-    AxisOption("Clip skip", int, apply_clip_skip, format_value_add_label, None),
-    AxisOption("Denoising", float, apply_field("denoising_strength"), format_value_add_label, None),
-    AxisOption("Hires upscaler", str, apply_field("hr_upscaler"), format_value_add_label, None),
-    AxisOption("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight"), format_value_add_label, None),
-    AxisOption("VAE", str, apply_vae, format_value_add_label, None),
-    AxisOption("Styles", str, apply_styles, format_value_add_label, None),
+    AxisOption("Nothing", str, do_nothing, format_value=format_nothing),
+    AxisOption("Seed", int, apply_field("seed")),
+    AxisOption("Var. seed", int, apply_field("subseed")),
+    AxisOption("Var. strength", float, apply_field("subseed_strength")),
+    AxisOption("Steps", int, apply_field("steps")),
+    AxisOption("CFG Scale", float, apply_field("cfg_scale")),
+    AxisOption("Prompt S/R", str, apply_prompt, format_value=format_value),
+    AxisOption("Prompt order", str_permutations, apply_order, format_value=format_value_join_list),
+    AxisOption("Sampler", str, apply_sampler, format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
+    AxisOption("Checkpoint name", str, apply_checkpoint, format_value=format_value, confirm=confirm_checkpoints, cost=1.0, choices=lambda: list(sd_models.checkpoints_list)),
+    AxisOption("Hypernetwork", str, apply_hypernetwork, format_value=format_value, confirm=confirm_hypernetworks, cost=0.2, choices=lambda: list(shared.hypernetworks)),
+    AxisOption("Hypernet str.", float, apply_hypernetwork_strength),
+    AxisOption("Sigma Churn", float, apply_field("s_churn")),
+    AxisOption("Sigma min", float, apply_field("s_tmin")),
+    AxisOption("Sigma max", float, apply_field("s_tmax")),
+    AxisOption("Sigma noise", float, apply_field("s_noise")),
+    AxisOption("Eta", float, apply_field("eta")),
+    AxisOption("Clip skip", int, apply_clip_skip),
+    AxisOption("Denoising", float, apply_field("denoising_strength")),
+    AxisOption("Hires upscaler", str, apply_field("hr_upscaler"), choices=lambda: [x.name for x in shared.sd_upscalers]),
+    AxisOption("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight")),
+    AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: list(sd_vae.vae_dict)),
+    AxisOption("Styles", str, apply_styles, choices=lambda: list(shared.prompt_styles.styles)),
 ]
 
 
-def draw_xy_grid(p, xs, ys, x_labels, y_labels, cell, draw_legend, include_lone_images):
+def draw_xy_grid(p, xs, ys, x_labels, y_labels, cell, draw_legend, include_lone_images, swap_axes_processing_order):
     ver_texts = [[images.GridAnnotation(y)] for y in y_labels]
     hor_texts = [[images.GridAnnotation(x)] for x in x_labels]
 
     # Temporary list of all the images that are generated to be populated into the grid.
     # Will be filled with empty images for any individual step that fails to process properly
-    image_cache = []
+    image_cache = [None] * (len(xs) * len(ys))
 
     processed_result = None
     cell_mode = "P"
-    cell_size = (1,1)
+    cell_size = (1, 1)
 
     state.job_count = len(xs) * len(ys) * p.n_iter
 
-    for iy, y in enumerate(ys):
+    def process_cell(x, y, ix, iy):
+        nonlocal image_cache, processed_result, cell_mode, cell_size
+
+        state.job = f"{ix + iy * len(xs) + 1} out of {len(xs) * len(ys)}"
+
+        processed: Processed = cell(x, y)
+
+        try:
+            # this dereference will throw an exception if the image was not processed
+            # (this happens in cases such as if the user stops the process from the UI)
+            processed_image = processed.images[0]
+
+            if processed_result is None:
+                # Use our first valid processed result as a template container to hold our full results
+                processed_result = copy(processed)
+                cell_mode = processed_image.mode
+                cell_size = processed_image.size
+                processed_result.images = [Image.new(cell_mode, cell_size)]
+
+            image_cache[ix + iy * len(xs)] = processed_image
+            if include_lone_images:
+                processed_result.images.append(processed_image)
+                processed_result.all_prompts.append(processed.prompt)
+                processed_result.all_seeds.append(processed.seed)
+                processed_result.infotexts.append(processed.infotexts[0])
+        except:
+            image_cache[ix + iy * len(xs)] = Image.new(cell_mode, cell_size)
+
+    if swap_axes_processing_order:
         for ix, x in enumerate(xs):
-            state.job = f"{ix + iy * len(xs) + 1} out of {len(xs) * len(ys)}"
-
-            processed:Processed = cell(x, y)
-            try:
-                # this dereference will throw an exception if the image was not processed
-                # (this happens in cases such as if the user stops the process from the UI)
-                processed_image = processed.images[0]
-                
-                if processed_result is None:
-                    # Use our first valid processed result as a template container to hold our full results
-                    processed_result = copy(processed)
-                    cell_mode = processed_image.mode
-                    cell_size = processed_image.size
-                    processed_result.images = [Image.new(cell_mode, cell_size)]
-
-                image_cache.append(processed_image)
-                if include_lone_images:
-                    processed_result.images.append(processed_image)
-                    processed_result.all_prompts.append(processed.prompt)
-                    processed_result.all_seeds.append(processed.seed)
-                    processed_result.infotexts.append(processed.infotexts[0])
-            except:
-                image_cache.append(Image.new(cell_mode, cell_size))
+            for iy, y in enumerate(ys):
+                process_cell(x, y, ix, iy)
+    else:
+        for iy, y in enumerate(ys):
+            for ix, x in enumerate(xs):
+                process_cell(x, y, ix, iy)
 
     if not processed_result:
         print("Unexpected error: draw_xy_grid failed to return even a single processed image")
-        return Processed()
+        return Processed(p, [])
 
     grid = images.image_grid(image_cache, rows=len(ys))
     if draw_legend:
@@ -266,12 +292,12 @@ class SharedSettingsStackHelper(object):
     def __enter__(self):
         self.CLIP_stop_at_last_layers = opts.CLIP_stop_at_last_layers
         self.hypernetwork = opts.sd_hypernetwork
-        self.model = shared.sd_model
         self.vae = opts.sd_vae
   
     def __exit__(self, exc_type, exc_value, tb):
-        modules.sd_models.reload_model_weights(self.model)
-        modules.sd_vae.reload_vae_weights(self.model, vae_file=find_vae(self.vae))
+        opts.data["sd_vae"] = self.vae
+        modules.sd_models.reload_model_weights()
+        modules.sd_vae.reload_vae_weights()
 
         hypernetwork.load_hypernetwork(self.hypernetwork)
         hypernetwork.apply_strength()
@@ -291,19 +317,45 @@ class Script(scripts.Script):
         return "X/Y plot"
 
     def ui(self, is_img2img):
-        current_axis_options = [x for x in axis_options if type(x) == AxisOption or type(x) == AxisOptionImg2Img and is_img2img]
+        current_axis_options = [x for x in axis_options if type(x) == AxisOption or x.is_img2img and is_img2img]
 
         with gr.Row():
-            x_type = gr.Dropdown(label="X type", choices=[x.label for x in current_axis_options], value=current_axis_options[1].label, type="index", elem_id=self.elem_id("x_type"))
-            x_values = gr.Textbox(label="X values", lines=1, elem_id=self.elem_id("x_values"))
+            with gr.Column(scale=19):
+                with gr.Row():
+                    x_type = gr.Dropdown(label="X type", choices=[x.label for x in current_axis_options], value=current_axis_options[1].label, type="index", elem_id=self.elem_id("x_type"))
+                    x_values = gr.Textbox(label="X values", lines=1, elem_id=self.elem_id("x_values"))
+                    fill_x_button = ToolButton(value=fill_values_symbol, elem_id="xy_grid_fill_x_tool_button", visible=False)
 
-        with gr.Row():
-            y_type = gr.Dropdown(label="Y type", choices=[x.label for x in current_axis_options], value=current_axis_options[0].label, type="index", elem_id=self.elem_id("y_type"))
-            y_values = gr.Textbox(label="Y values", lines=1, elem_id=self.elem_id("y_values"))
-        
-        draw_legend = gr.Checkbox(label='Draw legend', value=True, elem_id=self.elem_id("draw_legend"))
-        include_lone_images = gr.Checkbox(label='Include Separate Images', value=False, elem_id=self.elem_id("include_lone_images"))
-        no_fixed_seeds = gr.Checkbox(label='Keep -1 for seeds', value=False, elem_id=self.elem_id("no_fixed_seeds"))
+                with gr.Row():
+                    y_type = gr.Dropdown(label="Y type", choices=[x.label for x in current_axis_options], value=current_axis_options[0].label, type="index", elem_id=self.elem_id("y_type"))
+                    y_values = gr.Textbox(label="Y values", lines=1, elem_id=self.elem_id("y_values"))
+                    fill_y_button = ToolButton(value=fill_values_symbol, elem_id="xy_grid_fill_y_tool_button", visible=False)
+
+        with gr.Row(variant="compact"):
+            draw_legend = gr.Checkbox(label='Draw legend', value=True, elem_id=self.elem_id("draw_legend"))
+            include_lone_images = gr.Checkbox(label='Include Separate Images', value=False, elem_id=self.elem_id("include_lone_images"))
+            no_fixed_seeds = gr.Checkbox(label='Keep -1 for seeds', value=False, elem_id=self.elem_id("no_fixed_seeds"))
+            swap_axes_button = gr.Button(value="Swap axes", elem_id="xy_grid_swap_axes_button")
+
+        def swap_axes(x_type, x_values, y_type, y_values):
+            nonlocal current_axis_options
+            return current_axis_options[y_type].label, y_values, current_axis_options[x_type].label, x_values
+
+        swap_args = [x_type, x_values, y_type, y_values]
+        swap_axes_button.click(swap_axes, inputs=swap_args, outputs=swap_args)
+
+        def fill(x_type):
+            axis = axis_options[x_type]
+            return ", ".join(axis.choices()) if axis.choices else gr.update()
+
+        fill_x_button.click(fn=fill, inputs=[x_type], outputs=[x_values])
+        fill_y_button.click(fn=fill, inputs=[y_type], outputs=[y_values])
+
+        def select_axis(x_type):
+            return gr.Button.update(visible=axis_options[x_type].choices is not None)
+
+        x_type.change(fn=select_axis, inputs=[x_type], outputs=[fill_x_button])
+        y_type.change(fn=select_axis, inputs=[y_type], outputs=[fill_y_button])
 
         return [x_type, x_values, y_type, y_values, draw_legend, include_lone_images, no_fixed_seeds]
 
@@ -406,7 +458,15 @@ class Script(scripts.Script):
 
         grid_infotext = [None]
 
+        # If one of the axes is very slow to change between (like SD model
+        # checkpoint), then make sure it is in the outer iteration of the nested
+        # `for` loop.
+        swap_axes_processing_order = x_opt.cost > y_opt.cost
+
         def cell(x, y):
+            if shared.state.interrupted:
+                return Processed(p, [], p.seed, "")
+
             pc = copy(p)
             x_opt.apply(pc, x, xs)
             y_opt.apply(pc, y, ys)
@@ -441,7 +501,8 @@ class Script(scripts.Script):
                 y_labels=[y_opt.format_value(p, y_opt, y) for y in ys],
                 cell=cell,
                 draw_legend=draw_legend,
-                include_lone_images=include_lone_images
+                include_lone_images=include_lone_images,
+                swap_axes_processing_order=swap_axes_processing_order
             )
 
         if opts.grid_save: