wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-13 18:47:29 +08:00
Code Issues Actions 20 Packages Projects Releases Wiki Activity

prepare for multiple pin sets

Browse Source
This commit is contained in:
Rattus 2026-05-08 09:47:08 +10:00
parent 01eba77dba
commit eabf34c453
3 changed files with 28 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/model_management.py
View File

@ -1238,6 +1238,8 @@ def reset_cast_buffers():
model = loaded_model.model model = loaded_model.model
if model is not None and model.is_dynamic(): if model is not None and model.is_dynamic():
model.dynamic_pins[model.load_device]["active"] = False model.dynamic_pins[model.load_device]["active"] = False
model.partially_unload_ram(1e30, subsets=[ "patches" ])
model.dynamic_pins[model.load_device]["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0), [])
STREAM_CAST_BUFFERS.clear() STREAM_CAST_BUFFERS.clear()
STREAM_AIMDO_CAST_BUFFERS.clear() STREAM_AIMDO_CAST_BUFFERS.clear()

37
comfy/model_patcher.py
View File

@ -1508,8 +1508,8 @@ class ModelPatcherDynamic(ModelPatcher):
self.model.dynamic_pins = {} self.model.dynamic_pins = {}
if self.load_device not in self.model.dynamic_pins: if self.load_device not in self.model.dynamic_pins:
self.model.dynamic_pins[self.load_device] = { self.model.dynamic_pins[self.load_device] = {
"hostbuf": comfy_aimdo.host_buffer.HostBuffer(0), "weights": (comfy_aimdo.host_buffer.HostBuffer(0), []),
"stack": [], "patches": (comfy_aimdo.host_buffer.HostBuffer(0), []),
"failed": False, "failed": False,
"active": False, "active": False,
} }
@ -1703,25 +1703,26 @@ class ModelPatcherDynamic(ModelPatcher):
return freed return freed
def pinned_memory_size(self): def pinned_memory_size(self):
return self.model.dynamic_pins[self.load_device]["hostbuf"].size return (self.model.dynamic_pins[self.load_device]["weights"][0].size +
self.model.dynamic_pins[self.load_device]["patches"][0].size)
def partially_unload_ram(self, ram_to_unload): def partially_unload_ram(self, ram_to_unload, subsets=[ "weights", "patches" ]):
freed = 0 freed = 0
pin_state = self.model.dynamic_pins[self.load_device] pin_state = self.model.dynamic_pins[self.load_device]
hostbuf = pin_state["hostbuf"] for subset in subsets:
stack = pin_state["stack"] hostbuf, stack = pin_state[subset]
while len(stack) > 0: while len(stack) > 0:
module, offset = stack.pop() module, offset = stack.pop()
size = module._pin.numel() * module._pin.element_size() size = module._pin.numel() * module._pin.element_size()
del module._pin del module._pin
hostbuf.truncate(offset) hostbuf.truncate(offset)
comfy.model_management.TOTAL_PINNED_MEMORY -= size comfy.model_management.TOTAL_PINNED_MEMORY -= size
if comfy.model_management.TOTAL_PINNED_MEMORY < 0: if comfy.model_management.TOTAL_PINNED_MEMORY < 0:
comfy.model_management.TOTAL_PINNED_MEMORY = 0 comfy.model_management.TOTAL_PINNED_MEMORY = 0
freed += size freed += size
ram_to_unload -= size ram_to_unload -= size
if ram_to_unload <= 0: if ram_to_unload <= 0:
return freed return freed
return freed return freed
def patch_model(self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False): def patch_model(self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False):

13
comfy/pinned_memory.py
View File

@ -5,16 +5,17 @@ import comfy_aimdo.torch
from comfy.cli_args import args from comfy.cli_args import args
def get_pin(module): def get_pin(module, subset="weights"):
return getattr(module, "_pin", None) return getattr(module, "_pin", None)
def pin_memory(module): def pin_memory(module, subset="weights", size=None):
pin_state = module._pin_state pin_state = module._pin_state
if pin_state["failed"] or args.disable_pinned_memory or get_pin(module) is not None: if pin_state["failed"] or args.disable_pinned_memory or get_pin(module, subset) is not None:
return return
hostbuf = pin_state["hostbuf"] hostbuf, stack = pin_state[subset]
size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ]) if size is None:
size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
offset = hostbuf.size offset = hostbuf.size
comfy.model_management.ensure_pin_budget(size) comfy.model_management.ensure_pin_budget(size)
@ -26,6 +27,6 @@ def pin_memory(module):
module._pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)[offset:offset + size] module._pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)[offset:offset + size]
module._pin.untyped_storage()._comfy_hostbuf = hostbuf module._pin.untyped_storage()._comfy_hostbuf = hostbuf
pin_state["stack"].append((module, offset)) stack.append((module, offset))
comfy.model_management.TOTAL_PINNED_MEMORY += size comfy.model_management.TOTAL_PINNED_MEMORY += size
return True return True