From 41264bb3de13f631187597e27ceb172da9ea4871 Mon Sep 17 00:00:00 2001 From: James Walker Date: Thu, 6 Mar 2025 02:32:33 +0800 Subject: [PATCH 1/3] Use context manager to replace functions --- src/LanPaint/nodes.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py index 9d26b75..b510061 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 @@ -324,12 +325,17 @@ 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"] @@ -375,7 +381,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 +442,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 From 80a4e2b7a1aa04eba4799b255433e4412a3eeb5e Mon Sep 17 00:00:00 2001 From: scraed Date: Thu, 6 Mar 2025 09:32:09 +0800 Subject: [PATCH 2/3] add monkey patch of CFG guider --- src/LanPaint/nodes.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py index b510061..8231d7b 100644 --- a/src/LanPaint/nodes.py +++ b/src/LanPaint/nodes.py @@ -327,15 +327,23 @@ class KSAMPLER(comfy.samplers.KSAMPLER): @contextmanager def override_sample_function(): + original_outer_sample = comfy.samplers.CFGGuider.outer_sample + comfy.samplers.CFGGuider.outer_sample = CFGGuider_LanPaint.outer_sample + + original_predict_noise = comfy.samplers.CFGGuider.predict_noise + comfy.samplers.CFGGuider.predict_noise = CFGGuider_LanPaint.predict_noise + 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") + comfy.samplers.CFGGuider.predict_noise = original_predict_noise + comfy.samplers.CFGGuider.outer_sample = original_outer_sample + + KSAMPLER_NAMES = ["euler", "dpmpp_2m", "uni_pc"] From b14efd817060384ed278be83935151861dae5115 Mon Sep 17 00:00:00 2001 From: scraed Date: Thu, 6 Mar 2025 09:41:18 +0800 Subject: [PATCH 3/3] clean the previous monkey patch --- src/LanPaint/nodes.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py index 8231d7b..37bfbb7 100644 --- a/src/LanPaint/nodes.py +++ b/src/LanPaint/nodes.py @@ -7,8 +7,9 @@ import nodes import latent_preview from functools import partial from comfy.utils import repeat_to_batch_size +from comfy.samplers import * # Monkey patch comfy.samplers module by importing with absolute package path -exec(inspect.getsource(comfy.samplers).replace("from .", "from comfy.")) +#exec(inspect.getsource(comfy.samplers).replace("from .", "from comfy.")) def reshape_mask(input_mask, output_shape): dims = len(output_shape) - 2 @@ -66,8 +67,8 @@ class CFGGuider_LanPaint: def predict_noise(self, x, timestep, model_options={}, seed=None): return sampling_function_LanPaint(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, self.cfg_BIG, model_options=model_options, seed=seed) -CFGGuider.outer_sample = CFGGuider_LanPaint.outer_sample -CFGGuider.predict_noise = CFGGuider_LanPaint.predict_noise +#CFGGuider.outer_sample = CFGGuider_LanPaint.outer_sample +#CFGGuider.predict_noise = CFGGuider_LanPaint.predict_noise class KSamplerX0Inpaint: def __init__(self, model, sigmas):