add threshold for switchover from Karras to LDM noise schedule

This commit is contained in:
Lincoln Stein
2022-10-27 15:50:32 -04:00
parent 343ae8b7af
commit 1200fbd3bd
5 changed files with 45 additions and 18 deletions

View File

@ -10,7 +10,7 @@ from .shared_invokeai_diffusion import InvokeAIDiffuserComponent
# at this threshold, the scheduler will stop using the Karras
# noise schedule and start using the model's schedule
STEP_THRESHOLD = 30
STEP_THRESHOLD = 29
def cfg_apply_threshold(result, threshold = 0.0, scale = 0.7):
if threshold <= 0.0:
@ -68,6 +68,9 @@ class KSampler(Sampler):
self.sigmas = None
self.ds = None
self.s_in = None
self.karras_max = kwargs.get('karras_max',STEP_THRESHOLD)
if self.karras_max is None:
self.karras_max = STEP_THRESHOLD
def make_schedule(
self,
@ -97,11 +100,11 @@ class KSampler(Sampler):
device=self.device,
)
if ddim_num_steps >= STEP_THRESHOLD:
print(f'>> number of steps ({ddim_num_steps}) >= {STEP_THRESHOLD}: using model sigmas')
if ddim_num_steps >= self.karras_max:
print(f'>> Ksampler using model noise schedule (steps > {self.karras_max})')
self.sigmas = self.model_sigmas
else:
print(f'>> number of steps ({ddim_num_steps}) < {STEP_THRESHOLD}: using karras sigmas')
print(f'>> Ksampler using karras noise schedule (steps <= {self.karras_max})')
self.sigmas = self.karras_sigmas
# ALERT: We are completely overriding the sample() method in the base class, which