From 36250c78f5377f6ce4d02d6ceb49f7abe17f353a Mon Sep 17 00:00:00 2001
From: layerdiffusion <19834515+lllyasviel@users.noreply.github.com>
Date: Sun, 11 Aug 2024 18:07:50 -0700
Subject: [PATCH] fix flux img2img and highres.fix math

---
 modules/sd_samplers_kdiffusion.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index dedd6a4a..9292192f 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -139,7 +139,8 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
         sigma_sched = sigmas[steps - t_enc - 1:]
 
         x = x.to(noise)
-        xi = x + noise * sigma_sched[0]
+
+        xi = self.model_wrap.predictor.noise_scaling(sigma_sched[0], noise, x, max_denoise=False)
 
         if opts.img2img_extra_noise > 0:
             p.extra_generation_params["Extra noise"] = opts.img2img_extra_noise
@@ -198,9 +199,8 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
 
         if opts.sgm_noise_multiplier:
             p.extra_generation_params["SGM noise multiplier"] = True
-            x = x * torch.sqrt(1.0 + sigmas[0] ** 2.0)
-        else:
-            x = x * sigmas[0]
+
+        x = self.model_wrap.predictor.noise_scaling(sigmas[0], x, torch.zeros_like(x), max_denoise=opts.sgm_noise_multiplier)
 
         extra_params_kwargs = self.initialize(p)
         parameters = inspect.signature(self.func).parameters