Update sd_samplers.py

This commit is contained in:
hentailord85ez 2022-12-26 20:49:13 +00:00 committed by GitHub
parent 4af3ca5393
commit 4df5009acb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions

View File

@ -462,6 +462,9 @@ class KDiffusionSampler:
return extra_params_kwargs
def get_sigmas(self, p, steps):
disc = opts.always_discard_next_to_last_sigma or (self.config is not None and self.config.options.get('discard_next_to_last_sigma', False))
steps += 1 if disc else 0
if p.sampler_noise_scheduler_override:
sigmas = p.sampler_noise_scheduler_override(steps)
elif self.config is not None and self.config.options.get('scheduler', None) == 'karras':
@ -469,7 +472,7 @@ class KDiffusionSampler:
else:
sigmas = self.model_wrap.get_sigmas(steps)
if self.config is not None and self.config.options.get('discard_next_to_last_sigma', False):
if disc:
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
return sigmas