diff --git a/modules/processing.py b/modules/processing.py index 7eeb5191c..1da753a2c 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -79,7 +79,7 @@ class StableDiffusionProcessing: self.paste_to = None self.color_corrections = None self.denoising_strength: float = 0 - + self.sampler_noise_scheduler_override = None self.ddim_discretize = opts.ddim_discretize self.s_churn = opts.s_churn self.s_tmin = opts.s_tmin @@ -130,7 +130,7 @@ class Processed: self.s_tmin = p.s_tmin self.s_tmax = p.s_tmax self.s_noise = p.s_noise - + self.sampler_noise_scheduler_override = p.sampler_noise_scheduler_override self.prompt = self.prompt if type(self.prompt) != list else self.prompt[0] self.negative_prompt = self.negative_prompt if type(self.negative_prompt) != list else self.negative_prompt[0] self.seed = int(self.seed if type(self.seed) != list else self.seed[0]) diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index dff89c092..925222148 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -290,7 +290,10 @@ class KDiffusionSampler: def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None): steps, t_enc = setup_img2img_steps(p, steps) - sigmas = self.model_wrap.get_sigmas(steps) + if p.sampler_noise_scheduler_override: + sigmas = p.sampler_noise_scheduler_override(steps) + else: + sigmas = self.model_wrap.get_sigmas(steps) noise = noise * sigmas[steps - t_enc - 1] xi = x + noise @@ -306,7 +309,10 @@ class KDiffusionSampler: def sample(self, p, x, conditioning, unconditional_conditioning, steps=None): steps = steps or p.steps - sigmas = self.model_wrap.get_sigmas(steps) + if p.sampler_noise_scheduler_override: + sigmas = p.sampler_noise_scheduler_override(steps) + else: + sigmas = self.model_wrap.get_sigmas(steps) x = x * sigmas[0] extra_params_kwargs = self.initialize(p)