Merge pull request #6803 from space-nuko/xy-grid-performance-improvement

Optimize XY grid to run slower axes fewer times

scripts/xy_grid.py

@@ -178,76 +178,87 @@ 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"])
+AxisOption = namedtuple("AxisOption", ["label", "type", "apply", "format_value", "confirm", "cost"])
+AxisOptionImg2Img = namedtuple("AxisOptionImg2Img", ["label", "type", "apply", "format_value", "confirm", "cost"])
 
 
 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_nothing, None, 0),
+    AxisOption("Seed", int, apply_field("seed"), format_value_add_label, None, 0),
+    AxisOption("Var. seed", int, apply_field("subseed"), format_value_add_label, None, 0),
+    AxisOption("Var. strength", float, apply_field("subseed_strength"), format_value_add_label, None, 0),
+    AxisOption("Steps", int, apply_field("steps"), format_value_add_label, None, 0),
+    AxisOption("CFG Scale", float, apply_field("cfg_scale"), format_value_add_label, None, 0),
+    AxisOption("Prompt S/R", str, apply_prompt, format_value, None, 0),
+    AxisOption("Prompt order", str_permutations, apply_order, format_value_join_list, None, 0),
+    AxisOption("Sampler", str, apply_sampler, format_value, confirm_samplers, 0),
+    AxisOption("Checkpoint name", str, apply_checkpoint, format_value, confirm_checkpoints, 1.0),
+    AxisOption("Hypernetwork", str, apply_hypernetwork, format_value, confirm_hypernetworks, 0.2),
+    AxisOption("Hypernet str.", float, apply_hypernetwork_strength, format_value_add_label, None, 0),
+    AxisOption("Sigma Churn", float, apply_field("s_churn"), format_value_add_label, None, 0),
+    AxisOption("Sigma min", float, apply_field("s_tmin"), format_value_add_label, None, 0),
+    AxisOption("Sigma max", float, apply_field("s_tmax"), format_value_add_label, None, 0),
+    AxisOption("Sigma noise", float, apply_field("s_noise"), format_value_add_label, None, 0),
+    AxisOption("Eta", float, apply_field("eta"), format_value_add_label, None, 0),
+    AxisOption("Clip skip", int, apply_clip_skip, format_value_add_label, None, 0),
+    AxisOption("Denoising", float, apply_field("denoising_strength"), format_value_add_label, None, 0),
+    AxisOption("Hires upscaler", str, apply_field("hr_upscaler"), format_value_add_label, None, 0),
+    AxisOption("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight"), format_value_add_label, None, 0),
+    AxisOption("VAE", str, apply_vae, format_value_add_label, None, 0.7),
+    AxisOption("Styles", str, apply_styles, format_value_add_label, None, 0),
 ]
 
 
-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")
@@ -417,6 +428,11 @@ 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, "")
@@ -455,7 +471,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: