Merge pull request #12491 from AUTOMATIC1111/xyz-csv-and-dropdown-mode

Bring back CSV mode for XYZ grid

1 files changed, 78 insertions, 46 deletions

diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index d37b428f..da0e48aa 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -175,14 +175,22 @@ def do_nothing(p, x, xs):
 def format_nothing(p, opt, x):

     return ""

 

+

 def format_remove_path(p, opt, x):

     return os.path.basename(x)

 

+

 def str_permutations(x):

     """dummy function for specifying it in AxisOption's type when you want to get a list of permutations"""

     return x

 

 

+def list_to_csv_string(data_list):

+    with StringIO() as o:

+        csv.writer(o).writerow(data_list)

+        return o.getvalue().strip()

+

+

 class AxisOption:

     def __init__(self, label, type, apply, format_value=format_value_add_label, confirm=None, cost=0.0, choices=None):

         self.label = label

@@ -199,6 +207,7 @@ class AxisOptionImg2Img(AxisOption):
         super().__init__(*args, **kwargs)

         self.is_img2img = True

 

+

 class AxisOptionTxt2Img(AxisOption):

     def __init__(self, *args, **kwargs):

         super().__init__(*args, **kwargs)

@@ -286,11 +295,10 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
             cell_size = (processed_result.width, processed_result.height)

             if processed_result.images[0] is not None:

                 cell_mode = processed_result.images[0].mode

-                #This corrects size in case of batches:

+                # This corrects size in case of batches:

                 cell_size = processed_result.images[0].size

             processed_result.images[idx] = Image.new(cell_mode, cell_size)

 

-

     if first_axes_processed == 'x':

         for ix, x in enumerate(xs):

             if second_axes_processed == 'y':

@@ -348,9 +356,9 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
     if draw_legend:

         z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], title_texts, [[images.GridAnnotation()]])

     processed_result.images.insert(0, z_grid)

-    #TODO: Deeper aspects of the program rely on grid info being misaligned between metadata arrays, which is not ideal.

-    #processed_result.all_prompts.insert(0, processed_result.all_prompts[0])

-    #processed_result.all_seeds.insert(0, processed_result.all_seeds[0])

+    # TODO: Deeper aspects of the program rely on grid info being misaligned between metadata arrays, which is not ideal.

+    # processed_result.all_prompts.insert(0, processed_result.all_prompts[0])

+    # processed_result.all_seeds.insert(0, processed_result.all_seeds[0])

     processed_result.infotexts.insert(0, processed_result.infotexts[0])

 

     return processed_result

@@ -374,8 +382,8 @@ class SharedSettingsStackHelper(object):
 re_range = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\(([+-]\d+)\s*\))?\s*")

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

 

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

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

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

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

 

 

 class Script(scripts.Script):

@@ -390,19 +398,19 @@ 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"))

-                    x_values_dropdown = gr.Dropdown(label="X values",visible=False,multiselect=True,interactive=True)

+                    x_values_dropdown = gr.Dropdown(label="X values", visible=False, multiselect=True, interactive=True)

                     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"))

-                    y_values_dropdown = gr.Dropdown(label="Y values",visible=False,multiselect=True,interactive=True)

+                    y_values_dropdown = gr.Dropdown(label="Y values", visible=False, multiselect=True, interactive=True)

                     fill_y_button = ToolButton(value=fill_values_symbol, elem_id="xyz_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"))

-                    z_values_dropdown = gr.Dropdown(label="Z values",visible=False,multiselect=True,interactive=True)

+                    z_values_dropdown = gr.Dropdown(label="Z values", visible=False, multiselect=True, interactive=True)

                     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"):

@@ -414,6 +422,9 @@ class Script(scripts.Script):
                 include_sub_grids = gr.Checkbox(label='Include Sub Grids', value=False, elem_id=self.elem_id("include_sub_grids"))

             with gr.Column():

                 margin_size = gr.Slider(label="Grid margins (px)", minimum=0, maximum=500, value=0, step=2, elem_id=self.elem_id("margin_size"))

+            with gr.Column():

+                csv_mode = gr.Checkbox(label='Use text inputs instead of dropdowns', value=False, elem_id=self.elem_id("csv_mode"))

+

 

         with gr.Row(variant="compact", elem_id="swap_axes"):

             swap_xy_axes_button = gr.Button(value="Swap X/Y axes", elem_id="xy_grid_swap_axes_button")

@@ -430,50 +441,71 @@ class Script(scripts.Script):
         xz_swap_args = [x_type, x_values, x_values_dropdown, z_type, z_values, z_values_dropdown]

         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]

-            return axis.choices() if axis.choices else gr.update()

+        def fill(axis_type, csv_mode):

+            axis = self.current_axis_options[axis_type]

+            if axis.choices:

+                if csv_mode:

+                    return list_to_csv_string(axis.choices()), gr.update()

+                else:

+                    return gr.update(), axis.choices()

+            else:

+                return gr.update(), gr.update()

 

-        fill_x_button.click(fn=fill, inputs=[x_type], outputs=[x_values_dropdown])

-        fill_y_button.click(fn=fill, inputs=[y_type], outputs=[y_values_dropdown])

-        fill_z_button.click(fn=fill, inputs=[z_type], outputs=[z_values_dropdown])

+        fill_x_button.click(fn=fill, inputs=[x_type, csv_mode], outputs=[x_values, x_values_dropdown])

+        fill_y_button.click(fn=fill, inputs=[y_type, csv_mode], outputs=[y_values, y_values_dropdown])

+        fill_z_button.click(fn=fill, inputs=[z_type, csv_mode], outputs=[z_values, z_values_dropdown])

 

-        def select_axis(axis_type,axis_values_dropdown):

+        def select_axis(axis_type, axis_values, axis_values_dropdown, csv_mode):

             choices = self.current_axis_options[axis_type].choices

             has_choices = choices is not None

-            current_values = axis_values_dropdown

+

+            current_values = axis_values

+            current_dropdown_values = axis_values_dropdown

             if has_choices:

                 choices = choices()

-                if isinstance(current_values,str):

-                    current_values = current_values.split(",")

-                current_values = list(filter(lambda x: x in choices, current_values))

-            return gr.Button.update(visible=has_choices),gr.Textbox.update(visible=not has_choices),gr.update(choices=choices if has_choices else None,visible=has_choices,value=current_values)

+                if csv_mode:

+                    current_dropdown_values = list(filter(lambda x: x in choices, current_dropdown_values))

+                    current_values = list_to_csv_string(current_dropdown_values)

+                else:

+                    current_dropdown_values = [x.strip() for x in chain.from_iterable(csv.reader(StringIO(axis_values)))]

+                    current_dropdown_values = list(filter(lambda x: x in choices, current_dropdown_values))

+

+            return (gr.Button.update(visible=has_choices), gr.Textbox.update(visible=not has_choices or csv_mode, value=current_values),

+                    gr.update(choices=choices if has_choices else None, visible=has_choices and not csv_mode, value=current_dropdown_values))

+

+        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])

+        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])

+        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])

+

+        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):

+            _fill_x_button, _x_values, _x_values_dropdown = select_axis(x_type, x_values, x_values_dropdown, csv_mode)

+            _fill_y_button, _y_values, _y_values_dropdown = select_axis(y_type, y_values, y_values_dropdown, csv_mode)

+            _fill_z_button, _z_values, _z_values_dropdown = select_axis(z_type, z_values, z_values_dropdown, csv_mode)

+            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

 

-        x_type.change(fn=select_axis, inputs=[x_type,x_values_dropdown], outputs=[fill_x_button,x_values,x_values_dropdown])

-        y_type.change(fn=select_axis, inputs=[y_type,y_values_dropdown], outputs=[fill_y_button,y_values,y_values_dropdown])

-        z_type.change(fn=select_axis, inputs=[z_type,z_values_dropdown], outputs=[fill_z_button,z_values,z_values_dropdown])

+        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])

 

-        def get_dropdown_update_from_params(axis,params):

+        def get_dropdown_update_from_params(axis, params):

             val_key = f"{axis} Values"

-            vals = params.get(val_key,"")

+            vals = params.get(val_key, "")

             valslist = [x.strip() for x in chain.from_iterable(csv.reader(StringIO(vals))) if x]

-            return gr.update(value = valslist)

+            return gr.update(value=valslist)

 

         self.infotext_fields = (

             (x_type, "X Type"),

             (x_values, "X Values"),

-            (x_values_dropdown, lambda params:get_dropdown_update_from_params("X",params)),

+            (x_values_dropdown, lambda params: get_dropdown_update_from_params("X", params)),

             (y_type, "Y Type"),

             (y_values, "Y Values"),

-            (y_values_dropdown, lambda params:get_dropdown_update_from_params("Y",params)),

+            (y_values_dropdown, lambda params: get_dropdown_update_from_params("Y", params)),

             (z_type, "Z Type"),

             (z_values, "Z Values"),

-            (z_values_dropdown, lambda params:get_dropdown_update_from_params("Z",params)),

+            (z_values_dropdown, lambda params: get_dropdown_update_from_params("Z", params)),

         )

 

-        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, margin_size]

+        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, margin_size, csv_mode]

 

-    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, margin_size):

+    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, margin_size, csv_mode):

         if not no_fixed_seeds:

             modules.processing.fix_seed(p)

 

@@ -484,7 +516,7 @@ class Script(scripts.Script):
             if opt.label == 'Nothing':

                 return [0]

 

-            if opt.choices is not None:

+            if opt.choices is not None and not csv_mode:

                 valslist = vals_dropdown

             else:

                 valslist = [x.strip() for x in chain.from_iterable(csv.reader(StringIO(vals))) if x]

@@ -503,8 +535,8 @@ class Script(scripts.Script):
                         valslist_ext += list(range(start, end, step))

                     elif mc is not None:

                         start = int(mc.group(1))

-                        end   = int(mc.group(2))

-                        num   = int(mc.group(3)) if mc.group(3) is not None else 1

+                        end = int(mc.group(2))

+                        num = int(mc.group(3)) if mc.group(3) is not None else 1

 

                         valslist_ext += [int(x) for x in np.linspace(start=start, stop=end, num=num).tolist()]

                     else:

@@ -525,8 +557,8 @@ class Script(scripts.Script):
                         valslist_ext += np.arange(start, end + step, step).tolist()

                     elif mc is not None:

                         start = float(mc.group(1))

-                        end   = float(mc.group(2))

-                        num   = int(mc.group(3)) if mc.group(3) is not None else 1

+                        end = float(mc.group(2))

+                        num = int(mc.group(3)) if mc.group(3) is not None else 1

 

                         valslist_ext += np.linspace(start=start, stop=end, num=num).tolist()

                     else:

@@ -545,22 +577,22 @@ class Script(scripts.Script):
             return valslist

 

         x_opt = self.current_axis_options[x_type]

-        if x_opt.choices is not None:

-            x_values = ",".join(x_values_dropdown)

+        if x_opt.choices is not None and not csv_mode:

+            x_values = list_to_csv_string(x_values_dropdown)

         xs = process_axis(x_opt, x_values, x_values_dropdown)

 

         y_opt = self.current_axis_options[y_type]

-        if y_opt.choices is not None:

-            y_values = ",".join(y_values_dropdown)

+        if y_opt.choices is not None and not csv_mode:

+            y_values = list_to_csv_string(y_values_dropdown)

         ys = process_axis(y_opt, y_values, y_values_dropdown)

 

         z_opt = self.current_axis_options[z_type]

-        if z_opt.choices is not None:

-            z_values = ",".join(z_values_dropdown)

+        if z_opt.choices is not None and not csv_mode:

+            z_values = list_to_csv_string(z_values_dropdown)

         zs = process_axis(z_opt, z_values, z_values_dropdown)

 

         # this could be moved to common code, but unlikely to be ever triggered anywhere else

-        Image.MAX_IMAGE_PIXELS = None # disable check in Pillow and rely on check below to allow large custom image sizes

+        Image.MAX_IMAGE_PIXELS = None  # disable check in Pillow and rely on check below to allow large custom image sizes

         grid_mp = round(len(xs) * len(ys) * len(zs) * p.width * p.height / 1000000)

         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)'

 

@@ -720,7 +752,7 @@ class Script(scripts.Script):
             # Auto-save main and sub-grids:

             grid_count = z_count + 1 if z_count > 1 else 1

             for g in range(grid_count):

-                #TODO: See previous comment about intentional data misalignment.

+                # TODO: See previous comment about intentional data misalignment.

                 adj_g = g-1 if g > 0 else g

                 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)