From fa1baedad357c8f69910f83ead2e1b3ac85edc89 Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin" <git_commits@alexgoodwin.dev>
Date: Sun, 12 Feb 2023 23:58:18 -0800
Subject: [PATCH] further simplify #9

---
 scripts/dynamic_thresholding.py | 32 +++++++++++++++++++++++
 scripts/xyz_grid_support.py     | 45 ---------------------------------
 2 files changed, 32 insertions(+), 45 deletions(-)
 delete mode 100644 scripts/xyz_grid_support.py

diff --git a/scripts/dynamic_thresholding.py b/scripts/dynamic_thresholding.py
index d8eb295..9aecc74 100644
--- a/scripts/dynamic_thresholding.py
+++ b/scripts/dynamic_thresholding.py
@@ -241,3 +241,35 @@ class CustomCFGDenoiser(sd_samplers_kdiffusion.CFGDenoiser):
             actualRes = torch.einsum("laxy,ab -> lbxy", resRGB, coefs.inverse())
 
         return actualRes
+
+######################### XYZ Plot Script Support logic #########################
+
+def make_axis_options():
+    xyz_grid = [x for x in scripts.scripts_data if x.script_class.__module__ == "xyz_grid.py"][0].module
+    def apply_mimic_scale(p, x, xs):
+        if x != 0:
+            setattr(p, "dynthres_enabled", True)
+            setattr(p, "dynthres_mimic_scale", x)
+        else:
+            setattr(p, "dynthres_enabled", False)
+    def apply_scheduler(field):
+        def core(p, x, xs):
+            if x not in VALID_MODES:
+                raise RuntimeError(f"Unknown Scheduler: {x}")
+            setattr(p, field, x)
+        return core
+    extra_axis_options = [
+        xyz_grid.AxisOption("[DynThres] Mimic Scale", float, apply_mimic_scale),
+        xyz_grid.AxisOption("[DynThres] Threshold Percentile", float, xyz_grid.apply_field("dynthres_threshold_percentile")),
+        xyz_grid.AxisOption("[DynThres] Mimic Scheduler", str, apply_scheduler("dynthres_mimic_mode"), choices=lambda: VALID_MODES),
+        xyz_grid.AxisOption("[DynThres] Mimic minimum", float, xyz_grid.apply_field("dynthres_mimic_scale_min")),
+        xyz_grid.AxisOption("[DynThres] CFG Scheduler", str, apply_scheduler("dynthres_cfg_mode"), choices=lambda: VALID_MODES),
+        xyz_grid.AxisOption("[DynThres] CFG minimum", float, xyz_grid.apply_field("dynthres_cfg_scale_min")),
+        xyz_grid.AxisOption("[DynThres] Power value", float, xyz_grid.apply_field("dynthres_power_val"))
+    ]
+    xyz_grid.axis_options.extend(extra_axis_options)
+
+try:
+    make_axis_options()
+except Exception as e:
+    print(f"Failed to add support for X/Y/Z Plot Script because: {e}")
diff --git a/scripts/xyz_grid_support.py b/scripts/xyz_grid_support.py
deleted file mode 100644
index 3d1bbbf..0000000
--- a/scripts/xyz_grid_support.py
+++ /dev/null
@@ -1,45 +0,0 @@
-from pathlib import Path
-
-from modules import scripts
-from scripts.dynamic_thresholding import VALID_MODES
-
-
-def make_axis_options():
-    xyz_grid = [x for x in scripts.scripts_data if x.script_class.__module__ == "xyz_grid.py"][0].module
-
-    AxisOption = xyz_grid.AxisOption
-    apply_field = xyz_grid.apply_field
-
-    def apply_mimic_scale():
-        def core(p, x, xs):
-            if x != 0:
-                setattr(p, "dynthres_enabled", True)
-                setattr(p, "dynthres_mimic_scale", x)
-            else:
-                setattr(p, "dynthres_enabled", False)
-
-        return core
-
-    def apply_scheduler(field):
-        def core(p, x, xs):
-            if x not in VALID_MODES:
-                raise RuntimeError(f"Unknown Scheduler: {x}")
-
-            setattr(p, field, x)
-
-        return core
-
-    extra_axis_options = [
-        AxisOption("DT Mimic Scale", float, apply_mimic_scale()),
-        AxisOption("DT Threshold", float, apply_field("dynthres_threshold_percentile")),
-        AxisOption("DT Mimic Scheduler", str, apply_scheduler("dynthres_mimic_mode"), choices=lambda: VALID_MODES),
-        AxisOption("DT Mimic min", float, apply_field("dynthres_mimic_scale_min")),
-        AxisOption("DT CFG Scheduler", str, apply_scheduler("dynthres_cfg_mode"), choices=lambda: VALID_MODES),
-        AxisOption("DT CFG min", float, apply_field("dynthres_cfg_scale_min")),
-        AxisOption("DT Power val", float, apply_field("dynthres_power_val")),
-        AxisOption("DT Experiment", int, apply_field("dynthres_experiment_mode")),
-    ]
-
-    xyz_grid.axis_options.extend(extra_axis_options)
-
-make_axis_options()