Merge pull request #632 from nekoworkshop/dev

Adding ToMe parameters to available x/y/z plot options

modules/generation_parameters_copypaste.py

@@ -316,7 +316,7 @@ infotext_to_setting_name_mapping = [
     ('Token merging merge attention', 'token_merging_merge_attention'),
     ('Token merging merge cross attention', 'token_merging_merge_cross_attention'),
     ('Token merging merge mlp', 'token_merging_merge_mlp'),
-    ('Token merging maximum downsampling', 'token_merging_maximum_downsampling'),
+    ('Token merging maximum downsampling', 'token_merging_maximum_down_sampling'),
     ('Token merging stride x', 'token_merging_stride_x'),
     ('Token merging stride y', 'token_merging_stride_y')
 ]

modules/shared.py

@@ -428,14 +428,14 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
 
 options_templates.update(options_section(('token_merging', 'Token Merging'), {
     "token_merging": OptionInfo(False, "Enable redundant token merging via tomesd. This can provide significant speed and memory improvements.", gr.Checkbox),
-    "token_merging_ratio": OptionInfo(0.5, "Merging Ratio", gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}),
+    "token_merging_ratio": OptionInfo(0.5, "Merging Ratio. Higher merging ratio = faster generation, smaller VRAM usage, lower quality.", gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}),
     "token_merging_hr_only": OptionInfo(True, "Apply only to high-res fix pass. Disabling can yield a ~20-35% speedup on contemporary resolutions.", gr.Checkbox),
     "token_merging_ratio_hr": OptionInfo(0.5, "Merging Ratio (high-res pass) - If 'Apply only to high-res' is enabled, this will always be the ratio used.", gr.Slider, {"minimum": 0, "maximum": 0.9, "step": 0.1}),
     "token_merging_random": OptionInfo(False, "Use random perturbations - Can improve outputs for certain samplers. For others, it may cause visual artifacting.", gr.Checkbox),
-    "token_merging_merge_attention": OptionInfo(True, "Merge attention", gr.Checkbox),
-    "token_merging_merge_cross_attention": OptionInfo(False, "Merge cross attention", gr.Checkbox),
-    "token_merging_merge_mlp": OptionInfo(False, "Merge mlp", gr.Checkbox),
-    "token_merging_maximum_down_sampling": OptionInfo(1, "Maximum down sampling", gr.Dropdown, lambda: {"choices": ["1", "2", "4", "8"]}),
+    "token_merging_merge_attention": OptionInfo(True, "Merge attention (Recommend on)", gr.Checkbox),
+    "token_merging_merge_cross_attention": OptionInfo(False, "Merge cross attention (Recommend off)", gr.Checkbox),
+    "token_merging_merge_mlp": OptionInfo(False, "Merge mlp (Strongly recommend off)", gr.Checkbox),
+    "token_merging_maximum_down_sampling": OptionInfo(1, "Maximum down sampling", gr.Radio, lambda: {"choices": [1, 2, 4, 8]}),
     "token_merging_stride_x": OptionInfo(2, "Stride - X", gr.Slider, {"minimum": 2, "maximum": 8, "step": 2}),
     "token_merging_stride_y": OptionInfo(2, "Stride - Y", gr.Slider, {"minimum": 2, "maximum": 8, "step": 2})
 }))

scripts/xyz_grid.py

@@ -145,6 +145,15 @@ def apply_face_restore(p, opt, x):
 
     p.restore_faces = is_active
 
+def apply_token_merging_ratio_hr(p, x, xs):
+    opts.data["token_merging_ratio_hr"] = x
+
+def apply_token_merging_ratio(p, x, xs):
+    opts.data["token_merging_ratio"] = x
+
+def apply_token_merging_random(p, x, xs):
+    is_active = x.lower() in ('true', 'yes', 'y', '1')
+    opts.data["token_merging_random"] = is_active
 
 def format_value_add_label(p, opt, x):
     if type(x) == float:
@@ -226,6 +235,9 @@ axis_options = [
     AxisOption("Styles", str, apply_styles, choices=lambda: list(shared.prompt_styles.styles)),
     AxisOption("UniPC Order", int, apply_uni_pc_order, cost=0.5),
     AxisOption("Face restore", str, apply_face_restore, format_value=format_value),
+    AxisOption("ToMe ratio",float,apply_token_merging_ratio),
+    AxisOption("ToMe ratio for Hires fix",float,apply_token_merging_ratio_hr),
+    AxisOption("ToMe random pertubations",str,apply_token_merging_random, choices = lambda: ["Yes","No"])
 ]
 
 
@@ -342,11 +354,16 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
 
 class SharedSettingsStackHelper(object):
     def __enter__(self):
+        #Save overridden settings so they can be restored later.
         self.CLIP_stop_at_last_layers = opts.CLIP_stop_at_last_layers
         self.vae = opts.sd_vae
         self.uni_pc_order = opts.uni_pc_order
+        self.token_merging_ratio_hr = opts.token_merging_ratio_hr
+        self.token_merging_ratio = opts.token_merging_ratio
+        self.token_merging_random = opts.token_merging_random
 
     def __exit__(self, exc_type, exc_value, tb):
+        #Restore overriden settings after plot generation.
         opts.data["sd_vae"] = self.vae
         opts.data["uni_pc_order"] = self.uni_pc_order
         sd_models.reload_model_weights()
@@ -354,6 +371,9 @@ class SharedSettingsStackHelper(object):
 
         opts.data["CLIP_stop_at_last_layers"] = self.CLIP_stop_at_last_layers
 
+        opts.data["token_merging_ratio_hr"] = self.token_merging_ratio_hr
+        opts.data["token_merging_ratio"] = self.token_merging_ratio
+        opts.data["token_merging_random"] = self.token_merging_random
 
 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*")