From ce85639ba78ff7b7606fd5b31a6ffa944d463ec0 Mon Sep 17 00:00:00 2001
From: ThereforeGames <eric@sparknight.io>
Date: Wed, 5 Mar 2025 15:11:05 -0500
Subject: [PATCH] Update nodes.py

---
 src/LanPaint/nodes.py | 35 ++++++++++++++++++++++-------------
 1 file changed, 22 insertions(+), 13 deletions(-)

diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py
index 9d26b75..80b10f8 100644
--- a/src/LanPaint/nodes.py
+++ b/src/LanPaint/nodes.py
@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from inspect import cleandoc
 import inspect
 # import nodes.py
@@ -120,10 +121,10 @@ class KSamplerX0Inpaint:
                 x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size_i , current_times, sigma_x = self.sigma_x(abt), sigma_y = self.sigma_y(abt), args = args)  
             x = x_t #* ( 1+sigma**2 )**0.5
             # out is x_0
-            out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed)
+            out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
             out = out * denoise_mask + self.latent_image * latent_mask
         else:
-            out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed)
+            out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
         return out
     def mid_times(self, current_times, step_size):
         sigma, abt = current_times
@@ -142,7 +143,8 @@ class KSamplerX0Inpaint:
         lamb = self.chara_lamb
         beta = self.chara_beta * (1-abt)**0.5
 
-        x_0, x_0_BIG = self.inner_model(x_t, sigma, model_options=model_options, seed=seed)
+        x_0 = self.inner_model(x_t, sigma, model_options=model_options, seed=seed)
+        x_0_BIG = self.inner_model.cfg_BIG
         e_t = x_t / ((1 - abt) ** 0.5 * (1 + sigma**2) ** 0.5 )- (abt ** 0.5  / (1 - abt) ** 0.5) * x_0
         e_t_BIG = x_t / ((1 - abt) ** 0.5 * (1 + sigma**2) ** 0.5 )- (abt ** 0.5  / (1 - abt) ** 0.5) * x_0_BIG
         
@@ -324,15 +326,20 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
         samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
         return samples
 
-# Monkey patch ComfyUI's sample function to use our custom KSAMPLER
-comfy_sample_sample = inspect.getsource(comfy.sample.sample).replace("def sample(", "def comfy_sample_sample(").replace("comfy.samplers.KSampler", "KSampler")
-exec(comfy_sample_sample)
-# Modify nodes.common_ksampler function to use our patched sample function
-common_ksampler = inspect.getsource(nodes.common_ksampler).replace("comfy.sample.sample", "comfy_sample_sample")
-exec(common_ksampler)
+@contextmanager
+def override_sample_function():
+    original_sample = comfy.samplers.KSAMPLER.sample
+    comfy.samplers.KSAMPLER.sample = KSAMPLER.sample
+    print("comfy.samplers.KSAMPLER:sample function overriden")
+
+    try:
+        yield
+    finally:
+        comfy.samplers.KSAMPLER.sample = original_sample
+        print("comfy.samplers.KSAMPLER:sample function reverted")
 
 
-KSAMPLER_NAMES = ["euler", "dpmpp_2m", "uni_pc"]
+KSAMPLER_NAMES = ["euler", "lcm", "dpmpp_2m", "uni_pc"]
 
 class LanPaint_KSampler():
     @classmethod
@@ -375,7 +382,8 @@ class LanPaint_KSampler():
         model.LanPaint_StartSigma = 20.
         model.LanPaint_EndSigma = 1.
         model.LanPaint_cfg_BIG = cfg
-        return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
+        with override_sample_function():
+            return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
 class LanPaint_KSamplerAdvanced:
     @classmethod
     def INPUT_TYPES(s):
@@ -435,7 +443,8 @@ class LanPaint_KSamplerAdvanced:
         model.LanPaint_StartSigma = LanPaint_StartSigma
         model.LanPaint_EndSigma = LanPaint_EndSigma
         model.LanPaint_cfg_BIG = LanPaint_cfg_BIG
-        return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
+        with override_sample_function():
+            return nodes.common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
 
 
 # A dictionary that contains all nodes you want to export with their names
@@ -449,4 +458,4 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     "LanPaint_KSampler": "LanPaint KSampler",
     "LanPaint_KSamplerAdvanced": "LanPaint KSampler (Advanced)"
-}
+}
\ No newline at end of file