From 57987b58aacb2e52346860713b9a5b1d647bfbcd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phane=20du=20Hamel?= Date: Tue, 12 Aug 2025 13:18:49 +0200 Subject: [PATCH] update shift logic --- nodes.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/nodes.py b/nodes.py index d827637..ad96a29 100644 --- a/nodes.py +++ b/nodes.py @@ -68,17 +68,16 @@ def wan_ksampler(model_high_noise, model_low_noise, seed, steps, cfgs, sampler_n return (out, ) def set_shift(model,sigma_shift): - - sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow - sampling_type = comfy.model_sampling.CONST + model_sampling = model.get_model_object("model_sampling") + if not model_sampling: + sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow + sampling_type = comfy.model_sampling.CONST + class ModelSamplingAdvanced(sampling_base, sampling_type): + pass - class ModelSamplingAdvanced(sampling_base, sampling_type): - pass - - model_sampling = ModelSamplingAdvanced(model.model.model_config) + model_sampling = ModelSamplingAdvanced(model.model.model_config) model_sampling.set_parameters(shift=sigma_shift, multiplier=1000) model.add_object_patch("model_sampling", model_sampling) - model.add_object_patch("model_sampling", model_sampling) return model class WanMoeKSampler: