Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Bring back CSV mode for XYZ grid #12491

Merged
merged 9 commits into from
Aug 13, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 78 additions & 46 deletions scripts/xyz_grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -199,6 +207,7 @@ def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.is_img2img = True


class AxisOptionTxt2Img(AxisOption):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
Expand Down Expand Up @@ -286,11 +295,10 @@ def index(ix, iy, iz):
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':
Expand Down Expand Up @@ -348,9 +356,9 @@ def index(ix, iy, iz):
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
Expand All @@ -374,8 +382,8 @@ def __exit__(self, exc_type, exc_value, tb):
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):
Expand All @@ -390,19 +398,19 @@ def ui(self, is_img2img):
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"):
Expand All @@ -414,6 +422,9 @@ def ui(self, is_img2img):
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")
Expand All @@ -430,50 +441,71 @@ def swap_axes(axis1_type, axis1_values, axis1_values_dropdown, axis2_type, axis2
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)

Expand All @@ -484,7 +516,7 @@ def process_axis(opt, vals, vals_dropdown):
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]
Expand All @@ -503,8 +535,8 @@ def process_axis(opt, vals, vals_dropdown):
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:
Expand All @@ -525,8 +557,8 @@ def process_axis(opt, vals, vals_dropdown):
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:
Expand All @@ -545,22 +577,22 @@ def process_axis(opt, vals, vals_dropdown):
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)'

Expand Down Expand Up @@ -720,7 +752,7 @@ def cell(x, y, z, ix, iy, iz):
# 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)

Expand Down