wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-15 03:27:24 +08:00
Code Issues Actions 22 Packages Projects Releases Wiki Activity

pins: implement freeing intermediate for pinned memory

Browse Source 
Pinning is more important than inactive intermediates and the stream
pin buffer is more important than even active intermediates.
This commit is contained in:
Rattus 2026-05-13 09:17:23 +10:00
parent ee927aafa8
commit d61026d020
3 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
comfy/memory_management.py
View File

@ -157,7 +157,7 @@ def set_ram_cache_release_state(callback, headroom):
extra_ram_release_callback = callback extra_ram_release_callback = callback
RAM_CACHE_HEADROOM = max(0, int(headroom)) 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: if extra_ram_release_callback is None:
return 0 return 0
return extra_ram_release_callback(target) return extra_ram_release_callback(target, free_active=free_active)

2
comfy/model_management.py
View File

@ -1208,6 +1208,7 @@ def resize_pin_buffer(pin_buffer, size):
if size <= old_size: if size <= old_size:
return True return True
growth = size - old_size 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) ensure_pin_budget(growth, evict_active=True)
try: try:
pin_buffer.extend(size=size, reallocate=True) pin_buffer.extend(size=size, reallocate=True)
@ -1389,6 +1390,7 @@ def pin_memory(tensor):
return False return False
size = tensor.nbytes size = tensor.nbytes
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
ensure_pin_budget(size) ensure_pin_budget(size)
ptr = tensor.data_ptr() ptr = tensor.data_ptr()

1
comfy/pinned_memory.py
View File

@ -17,6 +17,7 @@ def pin_memory(module, subset="weights", size=None):
if size is None: if size is None:
size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ]) size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
offset = hostbuf.size offset = hostbuf.size
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
comfy.model_management.ensure_pin_budget(size) comfy.model_management.ensure_pin_budget(size)
try: try: