implements most of xyz grid script

scripts/xyz_grid.py

@@ -205,26 +205,30 @@ axis_options = [
 ]
 
 
-def draw_xy_grid(p, xs, ys, x_labels, y_labels, cell, draw_legend, include_lone_images, swap_axes_processing_order):
+def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_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]
+    title_texts = [[images.GridAnnotation(z)] for z in z_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 = [None] * (len(xs) * len(ys))
+    image_cache = [None] * (len(xs) * len(ys) * len(zs))
 
     processed_result = None
     cell_mode = "P"
     cell_size = (1, 1)
 
-    state.job_count = len(xs) * len(ys) * p.n_iter
+    state.job_count = len(xs) * len(ys) * len(zs) * p.n_iter
 
-    def process_cell(x, y, ix, iy):
+    def process_cell(x, y, z, ix, iy, iz):
         nonlocal image_cache, processed_result, cell_mode, cell_size
 
-        state.job = f"{ix + iy * len(xs) + 1} out of {len(xs) * len(ys)}"
+        def index(ix, iy, iz):
+            return ix + iy*len(xs) + iz*len(xs)*len(ys)
 
-        processed: Processed = cell(x, y)
+        state.job = f"{index(ix, iy, iz) + 1} out of {len(xs) * len(ys) * len(zs)}"
+
+        processed: Processed = cell(x, y, z)
 
         try:
             # this dereference will throw an exception if the image was not processed
@@ -238,33 +242,40 @@ def draw_xy_grid(p, xs, ys, x_labels, y_labels, cell, draw_legend, include_lone_
                 cell_size = processed_image.size
                 processed_result.images = [Image.new(cell_mode, cell_size)]
 
-            image_cache[ix + iy * len(xs)] = processed_image
+            image_cache[index(ix, iy, iz)] = 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)
+            image_cache[index(ix, iy, iz)] = Image.new(cell_mode, cell_size)
 
     if swap_axes_processing_order:
         for ix, x in enumerate(xs):
             for iy, y in enumerate(ys):
-                process_cell(x, y, ix, iy)
+                for iy, y in enumerate(zs):
+                    process_cell(x, y, z, ix, iy, iz)
     else:
         for iy, y in enumerate(ys):
             for ix, x in enumerate(xs):
-                process_cell(x, y, ix, iy)
+                for iz, z in enumerate(zs):
+                    process_cell(x, y, z, ix, iy, iz)
 
     if not processed_result:
-        print("Unexpected error: draw_xy_grid failed to return even a single processed image")
+        print("Unexpected error: draw_xyz_grid failed to return even a single processed image")
         return Processed(p, [])
 
-    grid = images.image_grid(image_cache, rows=len(ys))
+    for i, title_text in enumerate(title_texts):
-    if draw_legend:
+        start_index = i * len(xs) * len(ys)
-        grid = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts)
+        end_index = start_index + len(xs) * len(ys)
-
+        grid = images.image_grid(image_cache[start_index:end_index], rows=len(ys))
-    processed_result.images[0] = grid
+        if draw_legend:
+            grid = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts)
+        if i == 0: # First position is a placeholder as mentioned above, so it can be directly replaced
+            processed_result.images[0] = grid
+        else:
+            processed_result.images.insert(i, grid)
 
     return processed_result
 
@@ -291,7 +302,7 @@ re_range_count_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+
 
 class Script(scripts.Script):
     def title(self):
-        return "X/Y plot"
+        return "X/Y/Z plot"
 
     def ui(self, is_img2img):
         self.current_axis_options = [x for x in axis_options if type(x) == AxisOption or x.is_img2img == is_img2img]
@@ -301,24 +312,35 @@ class Script(scripts.Script):
                 with gr.Row():
                     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"))
                     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)
+                    fill_x_button = ToolButton(value=fill_values_symbol, elem_id="xyz_grid_fill_x_tool_button", visible=False)
 
                 with gr.Row():
                     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"))
                     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)
+                    fill_y_button = ToolButton(value=fill_values_symbol, elem_id="xyzz_grid_fill_y_tool_button", visible=False)
+
+                with gr.Row():
+                    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"))
+                    z_values = gr.Textbox(label="Z values", lines=1, elem_id=self.elem_id("z_values"))
+                    fill_z_button = ToolButton(value=fill_values_symbol, elem_id="xyz_grid_fill_z_tool_button", visible=False)
 
         with gr.Row(variant="compact", elem_id="axis_options"):
             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")
+            swap_xy_axes_button = gr.Button(value="Swap X/Y axes", elem_id="xy_grid_swap_axes_button")
+            swap_yz_axes_button = gr.Button(value="Swap Y/Z axes", elem_id="yz_grid_swap_axes_button")
+            swap_xz_axes_button = gr.Button(value="Swap X/Z axes", elem_id="xz_grid_swap_axes_button")
 
-        def swap_axes(x_type, x_values, y_type, y_values):
+        def swap_axes(axis1_type, axis1_values, axis2_type, axis2_values):
-            return self.current_axis_options[y_type].label, y_values, self.current_axis_options[x_type].label, x_values
+            return self.current_axis_options[axis2_type].label, axis2_values, self.current_axis_options[axis1_type].label, axis1_values
 
-        swap_args = [x_type, x_values, y_type, y_values]
+        xy_swap_args = [x_type, x_values, y_type, y_values]
-        swap_axes_button.click(swap_axes, inputs=swap_args, outputs=swap_args)
+        swap_xy_axes_button.click(swap_axes, inputs=xy_swap_args, outputs=xy_swap_args)
+        yz_swap_args = [y_type, y_values, z_type, z_values]
+        swap_yz_axes_button.click(swap_axes, inputs=yz_swap_args, outputs=yz_swap_args)
+        xz_swap_args = [x_type, x_values, z_type, z_values]
+        swap_xz_axes_button.click(swap_axes, inputs=xz_swap_args, outputs=xz_swap_args)
 
         def fill(x_type):
             axis = self.current_axis_options[x_type]
@@ -326,16 +348,18 @@ class Script(scripts.Script):
 
         fill_x_button.click(fn=fill, inputs=[x_type], outputs=[x_values])
         fill_y_button.click(fn=fill, inputs=[y_type], outputs=[y_values])
+        fill_z_button.click(fn=fill, inputs=[z_type], outputs=[z_values])
 
         def select_axis(x_type):
             return gr.Button.update(visible=self.current_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])
+        z_type.change(fn=select_axis, inputs=[z_type], outputs=[fill_z_button])
 
-        return [x_type, x_values, y_type, y_values, draw_legend, include_lone_images, no_fixed_seeds]
+        return [x_type, x_values, y_type, y_values, z_type, z_values, draw_legend, include_lone_images, no_fixed_seeds]
 
-    def run(self, p, x_type, x_values, y_type, y_values, draw_legend, include_lone_images, no_fixed_seeds):
+    def run(self, p, x_type, x_values, y_type, y_values, z_type, z_values, draw_legend, include_lone_images, no_fixed_seeds):
         if not no_fixed_seeds:
             modules.processing.fix_seed(p)
 
@@ -409,6 +433,9 @@ class Script(scripts.Script):
         y_opt = self.current_axis_options[y_type]
         ys = process_axis(y_opt, y_values)
 
+        z_opt = self.current_axis_options[z_type]
+        zs = process_axis(z_opt, z_values)
+
         def fix_axis_seeds(axis_opt, axis_list):
             if axis_opt.label in ['Seed', 'Var. seed']:
                 return [int(random.randrange(4294967294)) if val is None or val == '' or val == -1 else val for val in axis_list]
@@ -418,21 +445,26 @@ class Script(scripts.Script):
         if not no_fixed_seeds:
             xs = fix_axis_seeds(x_opt, xs)
             ys = fix_axis_seeds(y_opt, ys)
+            zs = fix_axis_seeds(z_opt, zs)
 
         if x_opt.label == 'Steps':
-            total_steps = sum(xs) * len(ys)
+            total_steps = sum(xs) * len(ys) * len(zs)
         elif y_opt.label == 'Steps':
-            total_steps = sum(ys) * len(xs)
+            total_steps = sum(ys) * len(xs) * len(zs)
+        elif z_opt.label == 'Steps':
+            total_steps = sum(zs) * len(xs) * len(ys)
         else:
-            total_steps = p.steps * len(xs) * len(ys)
+            total_steps = p.steps * len(xs) * len(ys) * len(zs)
 
         if isinstance(p, StableDiffusionProcessingTxt2Img) and p.enable_hr:
             if x_opt.label == "Hires steps":
-                total_steps += sum(xs) * len(ys)
+                total_steps += sum(xs) * len(ys) * len(zs)
             elif y_opt.label == "Hires steps":
-                total_steps += sum(ys) * len(xs)
+                total_steps += sum(ys) * len(xs) * len(zs)
+            elif z_opt.label == "Hires steps":
+                total_steps += sum(zs) * len(xs) * len(ys)
             elif p.hr_second_pass_steps:
-                total_steps += p.hr_second_pass_steps * len(xs) * len(ys)
+                total_steps += p.hr_second_pass_steps * len(xs) * len(ys) * len(zs)
             else:
                 total_steps *= 2
 
@@ -440,7 +472,8 @@ class Script(scripts.Script):
 
         image_cell_count = p.n_iter * p.batch_size
         cell_console_text = f"; {image_cell_count} images per cell" if image_cell_count > 1 else ""
-        print(f"X/Y plot will create {len(xs) * len(ys) * image_cell_count} images on a {len(xs)}x{len(ys)} grid{cell_console_text}. (Total steps to process: {total_steps})")
+        plural_s = 's' if len(zs) > 1 else ''
+        print(f"X/Y 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})")
         shared.total_tqdm.updateTotal(total_steps)
 
         grid_infotext = [None]
@@ -450,13 +483,14 @@ class Script(scripts.Script):
         # `for` loop.
         swap_axes_processing_order = x_opt.cost > y_opt.cost
 
-        def cell(x, y):
+        def cell(x, y, z):
             if shared.state.interrupted:
                 return Processed(p, [], p.seed, "")
 
             pc = copy(p)
             x_opt.apply(pc, x, xs)
             y_opt.apply(pc, y, ys)
+            z_opt.apply(pc, z, zs)
 
             res = process_images(pc)
 
@@ -475,17 +509,25 @@ class Script(scripts.Script):
                     if y_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:
                         pc.extra_generation_params["Fixed Y Values"] = ", ".join([str(y) for y in ys])
 
+                if z_opt.label != 'Nothing':
+                    pc.extra_generation_params["Z Type"] = z_opt.label
+                    pc.extra_generation_params["Z Values"] = z_values
+                    if z_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:
+                        pc.extra_generation_params["Fixed Z Values"] = ", ".join([str(z) for z in zs])
+
                 grid_infotext[0] = processing.create_infotext(pc, pc.all_prompts, pc.all_seeds, pc.all_subseeds)
 
             return res
 
         with SharedSettingsStackHelper():
-            processed = draw_xy_grid(
+            processed = draw_xyz_grid(
                 p,
                 xs=xs,
                 ys=ys,
+                zs=zs,
                 x_labels=[x_opt.format_value(p, x_opt, x) for x in xs],
                 y_labels=[y_opt.format_value(p, y_opt, y) for y in ys],
+                z_labels=[y_opt.format_value(p, z_opt, z) for z in zs],
                 cell=cell,
                 draw_legend=draw_legend,
                 include_lone_images=include_lone_images,
@@ -493,6 +535,6 @@ class Script(scripts.Script):
             )
 
         if opts.grid_save:
-            images.save_image(processed.images[0], p.outpath_grids, "xy_grid", info=grid_infotext[0], extension=opts.grid_format, prompt=p.prompt, seed=processed.seed, grid=True, p=p)
+            images.save_image(processed.images[0], p.outpath_grids, "xyz_grid", info=grid_infotext[0], extension=opts.grid_format, prompt=p.prompt, seed=processed.seed, grid=True, p=p)
 
         return processed