readd prior change but more safely this time i guess

aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
2023-12-07 03:17:07 -08:00
parent a139d3a23b
commit 249b8cb604
1 changed files with 2 additions and 0 deletions
@@ -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.