mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-20 12:39:05 +00:00
fce0398470
* pinned_memory: remove JIT RAM pressure release This doesn't work, as freeing intermediates for pins needs to be higher-priority than freeing pins-for-pins if and when you are going to do that. So this is too late as pins-for-pins is model load time and we dont have JIT pins-for-pins. * cacheing: Add a filter to only free intermediates from inactive wfs This is to get priorities in amongst pins straight. * mm: free inactive-ram from RAM cache first Stuff from inactive workflows should be freed before anything else. * caching: purge old ModelPatchers first Dont try and score them, just dump them at the first sign of trouble if they arent part of the workflow.
44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
import comfy.model_management
|
|
import comfy.memory_management
|
|
import comfy_aimdo.host_buffer
|
|
import comfy_aimdo.torch
|
|
|
|
from comfy.cli_args import args
|
|
|
|
def get_pin(module):
|
|
return getattr(module, "_pin", None)
|
|
|
|
def pin_memory(module):
|
|
if module.pin_failed or args.disable_pinned_memory or get_pin(module) is not None:
|
|
return
|
|
|
|
size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
|
|
|
|
if comfy.model_management.MAX_PINNED_MEMORY <= 0 or (comfy.model_management.TOTAL_PINNED_MEMORY + size) > comfy.model_management.MAX_PINNED_MEMORY:
|
|
module.pin_failed = True
|
|
return False
|
|
|
|
try:
|
|
hostbuf = comfy_aimdo.host_buffer.HostBuffer(size)
|
|
except RuntimeError:
|
|
module.pin_failed = True
|
|
return False
|
|
|
|
module._pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)
|
|
module._pin_hostbuf = hostbuf
|
|
comfy.model_management.TOTAL_PINNED_MEMORY += size
|
|
return True
|
|
|
|
def unpin_memory(module):
|
|
if get_pin(module) is None:
|
|
return 0
|
|
size = module._pin.numel() * module._pin.element_size()
|
|
|
|
comfy.model_management.TOTAL_PINNED_MEMORY -= size
|
|
if comfy.model_management.TOTAL_PINNED_MEMORY < 0:
|
|
comfy.model_management.TOTAL_PINNED_MEMORY = 0
|
|
|
|
del module._pin
|
|
del module._pin_hostbuf
|
|
return size
|