From 34b1f51f4a9b02899b4a96ce916b17063bca4b26 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 18 Aug 2025 14:45:21 -0700 Subject: [PATCH] Created Add Memory to Reserve node --- comfy/model_management.py | 11 ++++++++-- comfy/model_patcher.py | 19 ++++++++++++++++- comfy_extras/nodes_memory_reserve.py | 32 ++++++++++++++++++++++++++++ nodes.py | 1 + 4 files changed, 60 insertions(+), 3 deletions(-) create mode 100644 comfy_extras/nodes_memory_reserve.py diff --git a/comfy/model_management.py b/comfy/model_management.py index 2a9f18068..1f5b9ddb8 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -582,16 +582,23 @@ def free_memory(memory_required, device, keep_loaded=[]): soft_empty_cache() return unloaded_models +def get_models_memory_reserve(models): + total_reserve = 0 + for model in models: + total_reserve += model.get_model_memory_reserve(convert_to_bytes=True) + return total_reserve + def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False): cleanup_models_gc() global vram_state inference_memory = minimum_inference_memory() - extra_mem = max(inference_memory, memory_required + extra_reserved_memory()) + models_memory_reserve = get_models_memory_reserve(models) + extra_mem = max(inference_memory, memory_required + extra_reserved_memory() + models_memory_reserve) if minimum_memory_required is None: minimum_memory_required = extra_mem else: - minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_memory()) + minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_memory() + models_memory_reserve) models = set(models) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 52e76b5f3..f99b700ae 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -24,7 +24,7 @@ import inspect import logging import math import uuid -from typing import Callable, Optional +from typing import Callable, Optional, Union import torch @@ -84,6 +84,12 @@ def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_ model_options["disable_cfg1_optimization"] = True return model_options +def add_model_options_memory_reserve(model_options, memory_reserve_gb: float): + if "model_memory_reserve" not in model_options: + model_options["model_memory_reserve"] = [] + model_options["model_memory_reserve"].append(memory_reserve_gb) + return model_options + def create_model_options_clone(orig_model_options: dict): return comfy.patcher_extension.copy_nested_dicts(orig_model_options) @@ -439,6 +445,17 @@ class ModelPatcher: self.force_cast_weights = True self.patches_uuid = uuid.uuid4() #TODO: optimize by preventing a full model reload for this + def add_model_memory_reserve(self, memory_reserve_gb: float): + """Adds additional expected memory usage for the model, in gigabytes.""" + self.model_options = add_model_options_memory_reserve(self.model_options, memory_reserve_gb) + + def get_model_memory_reserve(self, convert_to_bytes: bool = False) -> Union[float, int]: + """Returns the total expected memory usage for the model in gigabytes, or bytes if convert_to_bytes is True.""" + total_reserve = sum(self.model_options.get("model_memory_reserve", [])) + if convert_to_bytes: + return total_reserve * 1024 * 1024 * 1024 + return total_reserve + def add_weight_wrapper(self, name, function): self.weight_wrapper_patches[name] = self.weight_wrapper_patches.get(name, []) + [function] self.patches_uuid = uuid.uuid4() diff --git a/comfy_extras/nodes_memory_reserve.py b/comfy_extras/nodes_memory_reserve.py new file mode 100644 index 000000000..3440287a2 --- /dev/null +++ b/comfy_extras/nodes_memory_reserve.py @@ -0,0 +1,32 @@ +from comfy_api.latest import io, ComfyExtension + +class MemoryReserveNode(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="AddMemoryToReserve", + display_name="Add Memory to Reserve", + description="Adds additional expected memory usage for the model, in gigabytes.", + category="advanced/debug/model", + inputs=[ + io.Model.Input("model", tooltip="The model to add memory reserve to."), + io.Float.Input("memory_reserve_gb", min=0.0, default=0.0, max=2048.0, step=0.1, tooltip="The additional expected memory usage for the model, in gigabytes."), + ], + outputs=[ + io.Model.Output(tooltip="The model with the additional memory reserve."), + ], + ) + + @classmethod + def execute(cls, model: io.Model.Type, memory_reserve_gb: float) -> io.NodeOutput: + model.add_model_memory_reserve(memory_reserve_gb) + return io.NodeOutput(model) + +class MemoryReserveExtension(ComfyExtension): + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + MemoryReserveNode, + ] + +def comfy_entrypoint(): + return MemoryReserveExtension() diff --git a/nodes.py b/nodes.py index 860a236aa..098ed4424 100644 --- a/nodes.py +++ b/nodes.py @@ -2321,6 +2321,7 @@ async def init_builtin_extra_nodes(): "nodes_edit_model.py", "nodes_tcfg.py", "nodes_context_windows.py", + "nodes_memory_reserve.py", ] import_failed = []