From 062483823738ed610d8d074ba63910c90e9d45b7 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 21 Mar 2024 14:49:11 -0400
Subject: [PATCH] Add inverse noise scaling function.

---
 comfy/model_sampling.py | 3 +++
 comfy/samplers.py       | 1 +
 2 files changed, 4 insertions(+)

diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py
index d325f76d..37976b32 100644
--- a/comfy/model_sampling.py
+++ b/comfy/model_sampling.py
@@ -20,6 +20,9 @@ class EPS:
         noise += latent_image
         return noise
 
+    def inverse_noise_scaling(self, sigma, latent):
+        return latent
+
 class V_PREDICTION(EPS):
     def calculate_denoised(self, sigma, model_output, model_input):
         sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
diff --git a/comfy/samplers.py b/comfy/samplers.py
index d721cb2e..3678dc81 100644
--- a/comfy/samplers.py
+++ b/comfy/samplers.py
@@ -546,6 +546,7 @@ class KSAMPLER(Sampler):
             k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
 
         samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
+        samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
         return samples