Merge branch 'master' into tensorboard

1 files changed, 80 insertions, 27 deletions

diff --git a/scripts/xy_grid.py b/scripts/xy_grid.py
index 5cca168a..f04d9b7e 100644
--- a/scripts/xy_grid.py
+++ b/scripts/xy_grid.py
@@ -10,13 +10,16 @@ import numpy as np
 import modules.scripts as scripts

 import gradio as gr

 

-from modules import images

+from modules import images, paths, sd_samplers, processing

 from modules.hypernetworks import hypernetwork

-from modules.processing import process_images, Processed, get_correct_sampler, StableDiffusionProcessingTxt2Img

+from modules.processing import process_images, Processed, StableDiffusionProcessingTxt2Img

 from modules.shared import opts, cmd_opts, state

 import modules.shared as shared

 import modules.sd_samplers

 import modules.sd_models

+import modules.sd_vae

+import glob

+import os

 import re

 

 

@@ -58,29 +61,19 @@ def apply_order(p, x, xs):
         prompt_tmp += part

         prompt_tmp += x[idx]

     p.prompt = prompt_tmp + p.prompt

-    

-

-def build_samplers_dict(p):

-    samplers_dict = {}

-    for i, sampler in enumerate(get_correct_sampler(p)):

-        samplers_dict[sampler.name.lower()] = i

-        for alias in sampler.aliases:

-            samplers_dict[alias.lower()] = i

-    return samplers_dict

 

 

 def apply_sampler(p, x, xs):

-    sampler_index = build_samplers_dict(p).get(x.lower(), None)

-    if sampler_index is None:

+    sampler_name = sd_samplers.samplers_map.get(x.lower(), None)

+    if sampler_name is None:

         raise RuntimeError(f"Unknown sampler: {x}")

 

-    p.sampler_index = sampler_index

+    p.sampler_name = sampler_name

 

 

 def confirm_samplers(p, xs):

-    samplers_dict = build_samplers_dict(p)

     for x in xs:

-        if x.lower() not in samplers_dict.keys():

+        if x.lower() not in sd_samplers.samplers_map:

             raise RuntimeError(f"Unknown sampler: {x}")

 

 

@@ -89,6 +82,7 @@ def apply_checkpoint(p, x, xs):
     if info is None:

         raise RuntimeError(f"Unknown checkpoint: {x}")

     modules.sd_models.reload_model_weights(shared.sd_model, info)

+    p.sd_model = shared.sd_model

 

 

 def confirm_checkpoints(p, xs):

@@ -123,6 +117,38 @@ def apply_clip_skip(p, x, xs):
     opts.data["CLIP_stop_at_last_layers"] = x

 

 

+def apply_upscale_latent_space(p, x, xs):

+    if x.lower().strip() != '0':

+        opts.data["use_scale_latent_for_hires_fix"] = True

+    else:

+        opts.data["use_scale_latent_for_hires_fix"] = False

+

+

+def find_vae(name: str):

+    if name.lower() in ['auto', 'none']:

+        return name

+    else:

+        vae_path = os.path.abspath(os.path.join(paths.models_path, 'VAE'))

+        found = glob.glob(os.path.join(vae_path, f'**/{name}.*pt'), recursive=True)

+        if found:

+            return found[0]

+        else:

+            return 'auto'

+

+

+def apply_vae(p, x, xs):

+    if x.lower().strip() == 'none':

+        modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file='None')

+    else:

+        found = find_vae(x)

+        if found:

+            v = modules.sd_vae.reload_vae_weights(shared.sd_model, vae_file=found)

+

+

+def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _):

+    p.styles = x.split(',')

+

+

 def format_value_add_label(p, opt, x):

     if type(x) == float:

         x = round(x, 8)

@@ -152,7 +178,6 @@ 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"])

 

@@ -177,6 +202,10 @@ axis_options = [
     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),

 ]

 

 

@@ -238,9 +267,11 @@ class SharedSettingsStackHelper(object):
         self.CLIP_stop_at_last_layers = opts.CLIP_stop_at_last_layers

         self.hypernetwork = opts.sd_hypernetwork

         self.model = shared.sd_model

+        self.vae = opts.sd_vae

   

     def __exit__(self, exc_type, exc_value, tb):

         modules.sd_models.reload_model_weights(self.model)

+        modules.sd_vae.reload_vae_weights(self.model, vae_file=find_vae(self.vae))

 

         hypernetwork.load_hypernetwork(self.hypernetwork)

         hypernetwork.apply_strength()

@@ -254,6 +285,7 @@ re_range_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+(?:.\d
 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):

     def title(self):

         return "X/Y plot"

@@ -262,16 +294,16 @@ class Script(scripts.Script):
         current_axis_options = [x for x in axis_options if type(x) == AxisOption or type(x) == AxisOptionImg2Img and is_img2img]

 

         with gr.Row():

-            x_type = gr.Dropdown(label="X type", choices=[x.label for x in current_axis_options], value=current_axis_options[1].label, visible=False, type="index", elem_id="x_type")

-            x_values = gr.Textbox(label="X values", visible=False, lines=1)

+            x_type = gr.Dropdown(label="X type", choices=[x.label for x in current_axis_options], value=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"))

 

         with gr.Row():

-            y_type = gr.Dropdown(label="Y type", choices=[x.label for x in current_axis_options], value=current_axis_options[0].label, visible=False, type="index", elem_id="y_type")

-            y_values = gr.Textbox(label="Y values", visible=False, lines=1)

+            y_type = gr.Dropdown(label="Y type", choices=[x.label for x in current_axis_options], value=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"))

         

-        draw_legend = gr.Checkbox(label='Draw legend', value=True)

-        include_lone_images = gr.Checkbox(label='Include Separate Images', value=False)

-        no_fixed_seeds = gr.Checkbox(label='Keep -1 for seeds', value=False)

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

 

         return [x_type, x_values, y_type, y_values, draw_legend, include_lone_images, no_fixed_seeds]

 

@@ -350,7 +382,7 @@ class Script(scripts.Script):
         ys = process_axis(y_opt, y_values)

 

         def fix_axis_seeds(axis_opt, axis_list):

-            if axis_opt.label in ['Seed','Var. seed']:

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

             else:

                 return axis_list

@@ -372,12 +404,33 @@ class Script(scripts.Script):
         print(f"X/Y plot will create {len(xs) * len(ys) * p.n_iter} images on a {len(xs)}x{len(ys)} grid. (Total steps to process: {total_steps * p.n_iter})")

         shared.total_tqdm.updateTotal(total_steps * p.n_iter)

 

+        grid_infotext = [None]

+

         def cell(x, y):

             pc = copy(p)

             x_opt.apply(pc, x, xs)

             y_opt.apply(pc, y, ys)

 

-            return process_images(pc)

+            res = process_images(pc)

+

+            if grid_infotext[0] is None:

+                pc.extra_generation_params = copy(pc.extra_generation_params)

+

+                if x_opt.label != 'Nothing':

+                    pc.extra_generation_params["X Type"] = x_opt.label

+                    pc.extra_generation_params["X Values"] = x_values

+                    if x_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:

+                        pc.extra_generation_params["Fixed X Values"] = ", ".join([str(x) for x in xs])

+

+                if y_opt.label != 'Nothing':

+                    pc.extra_generation_params["Y Type"] = y_opt.label

+                    pc.extra_generation_params["Y Values"] = y_values

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

+

+                grid_infotext[0] = processing.create_infotext(pc, pc.all_prompts, pc.all_seeds, pc.all_subseeds)

+

+            return res

 

         with SharedSettingsStackHelper():

             processed = draw_xy_grid(

@@ -392,6 +445,6 @@ class Script(scripts.Script):
             )

 

         if opts.grid_save:

-            images.save_image(processed.images[0], p.outpath_grids, "xy_grid", prompt=p.prompt, seed=processed.seed, grid=True, p=p)

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

 

         return processed