diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py index e7f8bc6a..b46202b7 100644 --- a/comfy/model_sampling.py +++ b/comfy/model_sampling.py @@ -11,6 +11,14 @@ class EPS: sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1)) return model_input - model_output * sigma + def noise_scaling(self, sigma, noise, latent_image, max_denoise=False): + if max_denoise: + noise = noise * torch.sqrt(1.0 + sigma ** 2.0) + else: + noise = noise * sigma + if latent_image is not None: + noise += latent_image + return noise class V_PREDICTION(EPS): def calculate_denoised(self, sigma, model_output, model_input): diff --git a/comfy/samplers.py b/comfy/samplers.py index e5569322..46d444f5 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -534,19 +534,13 @@ class KSAMPLER(Sampler): else: model_k.noise = noise - if self.max_denoise(model_wrap, sigmas): - noise = noise * torch.sqrt(1.0 + sigmas[0] ** 2.0) - else: - noise = noise * sigmas[0] + noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas)) k_callback = None total_steps = len(sigmas) - 1 if callback is not None: k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps) - if latent_image is not None: - noise += latent_image - samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options) return samples