diff options
Diffstat (limited to 'scripts/xyz_grid.py')
| -rw-r--r-- | scripts/xyz_grid.py | 177 |
1 files changed, 114 insertions, 63 deletions
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index 1010845e..939d8605 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -3,6 +3,7 @@ from copy import copy from itertools import permutations, chain
import random
import csv
+import os.path
from io import StringIO
from PIL import Image
import numpy as np
@@ -10,7 +11,7 @@ import numpy as np import modules.scripts as scripts
import gradio as gr
-from modules import images, sd_samplers, processing, sd_models, sd_vae, sd_samplers_kdiffusion
+from modules import images, sd_samplers, processing, sd_models, sd_vae, sd_samplers_kdiffusion, errors
from modules.processing import process_images, Processed, StableDiffusionProcessingTxt2Img
from modules.shared import opts, state
import modules.shared as shared
@@ -66,14 +67,6 @@ def apply_order(p, x, xs): p.prompt = prompt_tmp + p.prompt
-def apply_sampler(p, x, xs):
- sampler_name = sd_samplers.samplers_map.get(x.lower(), None)
- if sampler_name is None:
- raise RuntimeError(f"Unknown sampler: {x}")
-
- p.sampler_name = sampler_name
-
-
def confirm_samplers(p, xs):
for x in xs:
if x.lower() not in sd_samplers.samplers_map:
@@ -93,6 +86,15 @@ def confirm_checkpoints(p, xs): raise RuntimeError(f"Unknown checkpoint: {x}")
+def confirm_checkpoints_or_none(p, xs):
+ for x in xs:
+ if x in (None, "", "None", "none"):
+ continue
+
+ if modules.sd_models.get_closet_checkpoint_match(x) is None:
+ raise RuntimeError(f"Unknown checkpoint: {x}")
+
+
def apply_clip_skip(p, x, xs):
opts.data["CLIP_stop_at_last_layers"] = x
@@ -183,11 +185,25 @@ 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()
+
+
+def csv_string_to_list_strip(data_str):
+ return list(map(str.strip, chain.from_iterable(csv.reader(StringIO(data_str)))))
+
+
class AxisOption:
def __init__(self, label, type, apply, format_value=format_value_add_label, confirm=None, cost=0.0, choices=None):
self.label = label
@@ -204,6 +220,7 @@ class AxisOptionImg2Img(AxisOption): super().__init__(*args, **kwargs)
self.is_img2img = True
+
class AxisOptionTxt2Img(AxisOption):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@@ -221,9 +238,10 @@ axis_options = [ AxisOptionImg2Img("Image CFG Scale", float, apply_field("image_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),
- AxisOptionTxt2Img("Sampler", str, apply_sampler, format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
- AxisOptionImg2Img("Sampler", str, apply_sampler, format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]),
- AxisOption("Checkpoint name", str, apply_checkpoint, format_value=format_value, confirm=confirm_checkpoints, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list, key=str.casefold)),
+ AxisOptionTxt2Img("Sampler", str, apply_field("sampler_name"), format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers if x.name not in opts.hide_samplers]),
+ AxisOptionTxt2Img("Hires sampler", str, apply_field("hr_sampler_name"), confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img if x.name not in opts.hide_samplers]),
+ AxisOptionImg2Img("Sampler", str, apply_field("sampler_name"), format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img if x.name not in opts.hide_samplers]),
+ AxisOption("Checkpoint name", str, apply_checkpoint, format_value=format_remove_path, confirm=confirm_checkpoints, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list, key=str.casefold)),
AxisOption("Negative Guidance minimum sigma", float, apply_field("s_min_uncond")),
AxisOption("Sigma Churn", float, apply_field("s_churn")),
AxisOption("Sigma min", float, apply_field("s_tmin")),
@@ -236,6 +254,8 @@ axis_options = [ AxisOption("Eta", float, apply_field("eta")),
AxisOption("Clip skip", int, apply_clip_skip),
AxisOption("Denoising", float, apply_field("denoising_strength")),
+ AxisOption("Initial noise multiplier", float, apply_field("initial_noise_multiplier")),
+ AxisOption("Extra noise", float, apply_override("img2img_extra_noise")),
AxisOptionTxt2Img("Hires upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]),
AxisOptionImg2Img("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight")),
AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)),
@@ -245,6 +265,10 @@ axis_options = [ AxisOption("Token merging ratio", float, apply_override('token_merging_ratio')),
AxisOption("Token merging ratio high-res", float, apply_override('token_merging_ratio_hr')),
AxisOption("Always discard next-to-last sigma", str, apply_override('always_discard_next_to_last_sigma', boolean=True), choices=boolean_choice(reverse=True)),
+ AxisOption("SGM noise multiplier", str, apply_override('sgm_noise_multiplier', boolean=True), choices=boolean_choice(reverse=True)),
+ AxisOption("Refiner checkpoint", str, apply_field('refiner_checkpoint'), format_value=format_remove_path, confirm=confirm_checkpoints_or_none, cost=1.0, choices=lambda: ['None'] + sorted(sd_models.checkpoints_list, key=str.casefold)),
+ AxisOption("Refiner switch at", float, apply_field('refiner_switch_at')),
+ AxisOption("RNG source", str, apply_override("randn_source"), choices=lambda: ["GPU", "CPU", "NV"]),
]
@@ -290,11 +314,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':
@@ -352,9 +375,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
@@ -378,8 +401,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):
@@ -394,19 +417,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"):
@@ -418,6 +441,8 @@ 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")
@@ -434,50 +459,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
+
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)
-
- 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])
-
- def get_dropdown_update_from_params(axis,params):
+ if csv_mode:
+ if axis_values_dropdown:
+ axis_values = list_to_csv_string(list(filter(lambda x: x in choices, axis_values_dropdown)))
+ axis_values_dropdown = []
+ else:
+ if axis_values:
+ axis_values_dropdown = list(filter(lambda x: x in choices, csv_string_to_list_strip(axis_values)))
+ axis_values = ""
+
+ return (gr.Button.update(visible=has_choices), gr.Textbox.update(visible=not has_choices or csv_mode, value=axis_values),
+ gr.update(choices=choices if has_choices else None, visible=has_choices and not csv_mode, value=axis_values_dropdown))
+
+ 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
+
+ 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):
val_key = f"{axis} Values"
- 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)
+ vals = params.get(val_key, "")
+ valslist = csv_string_to_list_strip(vals)
+ 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)
@@ -488,10 +534,10 @@ 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]
+ valslist = csv_string_to_list_strip(vals)
if opt.type == int:
valslist_ext = []
@@ -507,8 +553,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:
@@ -529,8 +575,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:
@@ -549,22 +595,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)'
@@ -648,7 +694,12 @@ class Script(scripts.Script): y_opt.apply(pc, y, ys)
z_opt.apply(pc, z, zs)
- res = process_images(pc)
+ try:
+ res = process_images(pc)
+ except Exception as e:
+ errors.display(e, "generating image for xyz plot")
+
+ res = Processed(p, [], p.seed, "")
# Sets subgrid infotexts
subgrid_index = 1 + iz
@@ -719,7 +770,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)
|