From 58152b5606edd5042fea37bfcd160acb5e27bf76 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 5 Aug 2025 00:54:22 -0700 Subject: [PATCH] Add prepare_sampling wrapper for context window to more accurately estimate latent memory requirements, fixed merging wrappers/callbacks dicts in prepare_model_patcher --- comfy/sampler_helpers.py | 6 +++--- comfy_extras/nodes_context_windows.py | 22 ++++++++++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py index 8dbc41455..c6f88ef2f 100644 --- a/comfy/sampler_helpers.py +++ b/comfy/sampler_helpers.py @@ -149,7 +149,7 @@ def cleanup_models(conds, models): cleanup_additional_models(set(control_cleanup)) -def prepare_model_patcher(model: 'ModelPatcher', conds, model_options: dict): +def prepare_model_patcher(model: ModelPatcher, conds, model_options: dict): ''' Registers hooks from conds. ''' @@ -158,8 +158,8 @@ def prepare_model_patcher(model: 'ModelPatcher', conds, model_options: dict): for k in conds: get_hooks_from_cond(conds[k], hooks) # add wrappers and callbacks from ModelPatcher to transformer_options - model_options["transformer_options"]["wrappers"] = comfy.patcher_extension.copy_nested_dicts(model.wrappers) - model_options["transformer_options"]["callbacks"] = comfy.patcher_extension.copy_nested_dicts(model.callbacks) + comfy.patcher_extension.merge_nested_dicts(model_options["transformer_options"].setdefault("wrappers", {}), model.wrappers) + comfy.patcher_extension.merge_nested_dicts(model_options["transformer_options"].setdefault("callbacks", {}), model.callbacks) # begin registering hooks registered = comfy.hooks.HookGroup() target_dict = comfy.hooks.create_target_dict(comfy.hooks.EnumWeightTarget.Model) diff --git a/comfy_extras/nodes_context_windows.py b/comfy_extras/nodes_context_windows.py index 9be5fc43d..91823e620 100644 --- a/comfy_extras/nodes_context_windows.py +++ b/comfy_extras/nodes_context_windows.py @@ -1,6 +1,27 @@ from __future__ import annotations from comfy_api.latest import ComfyExtension, io import comfy.context_windows +import comfy.patcher_extension +import torch + + +def _prepare_sampling_wrapper(executor, model, noise_shape: torch.Tensor, *args, **kwargs): + # TODO: handle various dims instead of defaulting to 0th + # limit noise_shape length to context_length for more accurate vram use estimation + model_options = kwargs.get("model_options", None) + if model_options is None: + raise Exception("model_options not found in prepare_sampling_wrapper; this should never happen, something went wrong.") + handler: comfy.context_windows.IndexListContextHandler = model_options.get("context_handler", None) + if handler is not None: + noise_shape = [min(noise_shape[0], handler.context_length)] + list(noise_shape[1:]) + return executor(model, noise_shape, *args, **kwargs) + + +def create_prepare_sampling_wrapper(model_options: dict): + comfy.patcher_extension.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.PREPARE_SAMPLING, + "ContextWindows_prepare_sampling", + _prepare_sampling_wrapper, + model_options, is_model_options=True) class ContextWindowsNode(io.ComfyNode): @@ -37,6 +58,7 @@ class ContextWindowsNode(io.ComfyNode): context_length=context_length, context_overlap=context_overlap, dim=dim) + create_prepare_sampling_wrapper(model.model_options) return io.NodeOutput(model)