From e346d8584e30996455afcc3773f16442f24c3679 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 9 Apr 2025 21:43:35 +0800 Subject: [PATCH] Add prepare_sampling wrapper allowing custom nodes to more accurately report noise_shape (#7500) --- comfy/patcher_extension.py | 1 + comfy/sampler_helpers.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/comfy/patcher_extension.py b/comfy/patcher_extension.py index 859758244..965958f4c 100644 --- a/comfy/patcher_extension.py +++ b/comfy/patcher_extension.py @@ -48,6 +48,7 @@ def get_all_callbacks(call_type: str, transformer_options: dict, is_model_option class WrappersMP: OUTER_SAMPLE = "outer_sample" + PREPARE_SAMPLING = "prepare_sampling" SAMPLER_SAMPLE = "sampler_sample" CALC_COND_BATCH = "calc_cond_batch" APPLY_MODEL = "apply_model" diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py index 92ec7ca7a..96a3040a1 100644 --- a/comfy/sampler_helpers.py +++ b/comfy/sampler_helpers.py @@ -106,6 +106,13 @@ def cleanup_additional_models(models): def prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None): + executor = comfy.patcher_extension.WrapperExecutor.new_executor( + _prepare_sampling, + comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.PREPARE_SAMPLING, model_options, is_model_options=True) + ) + return executor.execute(model, noise_shape, conds, model_options=model_options) + +def _prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None): real_model: BaseModel = None models, inference_memory = get_additional_models(conds, model.model_dtype()) models += get_additional_models_from_model_options(model_options)