Move noise scaling to object with sampling math.

This commit is contained in:
comfyanonymous 2024-03-01 12:54:38 -05:00
parent cb7c3a2921
commit c62e836167
2 changed files with 9 additions and 7 deletions

View File

@ -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):

View File

@ -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