From 3e48b9ff85fc9cce5010fa733969a70341a7cf66 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 22 Oct 2022 23:02:50 -0400 Subject: [PATCH] cut over from karras to model noise schedule for higher steps The k_samplers come with a "karras" noise schedule which performs very well at low step counts but becomes noisy at higher ones. This commit introduces a threshold (currently 30 steps) at which the k samplers will switch over from using karras to the older model noise schedule. --- ldm/models/diffusion/ksampler.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py index ac0615b30c..677e051be7 100644 --- a/ldm/models/diffusion/ksampler.py +++ b/ldm/models/diffusion/ksampler.py @@ -12,6 +12,10 @@ from ldm.modules.diffusionmodules.util import ( extract_into_tensor, ) +# at this threshold, the scheduler will stop using the Karras +# noise schedule and start using the model's schedule +STEP_THRESHOLD = 30 + def cfg_apply_threshold(result, threshold = 0.0, scale = 0.7): if threshold <= 0.0: return result @@ -98,8 +102,13 @@ class KSampler(Sampler): rho=7., device=self.device, ) - self.sigmas = self.model_sigmas - #self.sigmas = self.karras_sigmas + + if ddim_num_steps >= STEP_THRESHOLD: + print(f'>> number of steps ({ddim_num_steps}) >= {STEP_THRESHOLD}: using model sigmas') + self.sigmas = self.model_sigmas + else: + print(f'>> number of steps ({ddim_num_steps}) < {STEP_THRESHOLD}: using karras sigmas') + self.sigmas = self.karras_sigmas # ALERT: We are completely overriding the sample() method in the base class, which # means that inpainting will not work. To get this to work we need to be able to