From 249b8cb6048f18e569873d28f235e233e88c52a4 Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin" <git_commits@alexgoodwin.dev>
Date: Thu, 7 Dec 2023 03:17:07 -0800
Subject: [PATCH] readd prior change but more safely this time i guess

aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
---
 scripts/dynamic_thresholding.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/scripts/dynamic_thresholding.py b/scripts/dynamic_thresholding.py
index 07d80fa..082fe46 100644
--- a/scripts/dynamic_thresholding.py
+++ b/scripts/dynamic_thresholding.py
@@ -214,6 +214,8 @@ class CustomCFGDenoiser(cfgdenoisekdiff):
         weights = torch.tensor(conds_list, device=uncond.device).select(2, 1)
         weights = weights.reshape(*weights.shape, 1, 1, 1)
         self.main_class.step = self.step
+        if hasattr(self, 'total_steps'):
+            self.main_class.max_steps = self.total_steps
 
         if self.main_class.experiment_mode >= 4 and self.main_class.experiment_mode <= 5:
             # https://arxiv.org/pdf/2305.08891.pdf "Rescale CFG". It's not good, but if you want to test it, just set experiment_mode = 4 + phi.