mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-16 01:00:49 +08:00
439c178c2c
Add a python for managing pinned memory of the weight/bias module level. This allocates, pins and attached a tensor to a module for the pin for this module. It does not set the weight, just allocates a singular ram buffer for population and bulk DMA transfer.
35 lines
970 B
Python
35 lines
970 B
Python
import torch
|
|
import logging
|
|
import comfy.model_management
|
|
import comfy.memory_management
|
|
|
|
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
|
|
#FIXME: This is a RAM cache trigger event
|
|
params = [ module.weight, module.bias ]
|
|
size = comfy.memory_management.vram_aligned_size(params)
|
|
try:
|
|
pin = torch.empty((size,), dtype=torch.uint8)
|
|
if comfy.model_management.pin_memory(pin):
|
|
module._pin = pin
|
|
else:
|
|
module.pin_failed = True
|
|
return False
|
|
except:
|
|
module.pin_failed = True
|
|
return False
|
|
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.unpin_memory(module._pin)
|
|
return size
|