Skip to content

Commit 0a29fde

Browse files
committed
use total_steps if available, for #83
1 parent 2998d4c commit 0a29fde

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

scripts/dynamic_thresholding.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,8 @@ def combine_denoised(self, x_out, conds_list, uncond, cond_scale):
212212
weights = torch.tensor(conds_list, device=uncond.device).select(2, 1)
213213
weights = weights.reshape(*weights.shape, 1, 1, 1)
214214
self.main_class.step = self.step
215+
if self.total_steps:
216+
self.main_class.max_steps = self.total_steps
215217

216218
if self.main_class.experiment_mode >= 4 and self.main_class.experiment_mode <= 5:
217219
# 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.

0 commit comments

Comments
 (0)