diff --git a/comfy/memory_management.py b/comfy/memory_management.py index 4a628b05c..7645064f5 100644 --- a/comfy/memory_management.py +++ b/comfy/memory_management.py @@ -157,7 +157,7 @@ def set_ram_cache_release_state(callback, headroom): extra_ram_release_callback = callback RAM_CACHE_HEADROOM = max(0, int(headroom)) -def extra_ram_release(target): +def extra_ram_release(target, free_active=False): if extra_ram_release_callback is None: return 0 - return extra_ram_release_callback(target) + return extra_ram_release_callback(target, free_active=free_active) diff --git a/comfy/model_management.py b/comfy/model_management.py index c1d0901fc..697359d3a 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1208,6 +1208,7 @@ def resize_pin_buffer(pin_buffer, size): if size <= old_size: return True growth = size - old_size + comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM, free_active=True) ensure_pin_budget(growth, evict_active=True) try: pin_buffer.extend(size=size, reallocate=True) @@ -1389,6 +1390,7 @@ def pin_memory(tensor): return False size = tensor.nbytes + comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM) ensure_pin_budget(size) ptr = tensor.data_ptr() diff --git a/comfy/pinned_memory.py b/comfy/pinned_memory.py index 208c777f8..35cbbcd9e 100644 --- a/comfy/pinned_memory.py +++ b/comfy/pinned_memory.py @@ -17,6 +17,7 @@ def pin_memory(module, subset="weights", size=None): if size is None: size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ]) offset = hostbuf.size + comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM) comfy.model_management.ensure_pin_budget(size) try: