wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-26 09:49:26 +08:00
Code Issues Actions 25 Packages Projects Releases Wiki Activity

Merge branch 'master' into feat/api-nodes/new-gemini-text-node

Browse Source
This commit is contained in:
Alexander Piskun 2026-06-05 18:47:52 +03:00 committed by GitHub
commit bf20c6530e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 21 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/model_management.py
View File

@ -651,8 +651,7 @@ def ensure_pin_budget(size, evict_active=False):
to_free = shortfall + PIN_PRESSURE_HYSTERESIS to_free = shortfall + PIN_PRESSURE_HYSTERESIS
return free_pins(to_free, evict_active=evict_active) >= shortfall return free_pins(to_free, evict_active=evict_active) >= shortfall
def ensure_pin_registerable(size, evict_active=True): def free_registrations(shortfall, evict_active=True):
shortfall = TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY
if MAX_PINNED_MEMORY <= 0: if MAX_PINNED_MEMORY <= 0:
return False return False
if shortfall <= 0: if shortfall <= 0:
@ -674,6 +673,9 @@ def ensure_pin_registerable(size, evict_active=True):
return True return True
return shortfall <= REGISTERABLE_PIN_HYSTERESIS return shortfall <= REGISTERABLE_PIN_HYSTERESIS
def ensure_pin_registerable(size, evict_active=True):
return free_registrations(TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY, evict_active=evict_active)
class LoadedModel: class LoadedModel:
def __init__(self, model: ModelPatcher): def __init__(self, model: ModelPatcher):
self._set_model(model) self._set_model(model)

19
comfy/pinned_memory.py
View File

@ -89,13 +89,26 @@ def pin_memory(module, subset="weights", size=None):
not comfy.model_management.ensure_pin_registerable(registerable_size)): not comfy.model_management.ensure_pin_registerable(registerable_size)):
return _steal_pin(module, stack, buckets, size, priority) return _steal_pin(module, stack, buckets, size, priority)
extended = False
try: try:
hostbuf.extend(size=size) hostbuf.extend(size=size, register=False)
extended = True
pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)[offset:offset + size]
pin.untyped_storage()._comfy_hostbuf = hostbuf
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
comfy.model_management.free_registrations(size)
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
del pin
hostbuf.truncate(offset, do_unregister=False)
return _steal_pin(module, stack, buckets, size, priority)
except RuntimeError: except RuntimeError:
if extended:
hostbuf.truncate(offset, do_unregister=False)
return _steal_pin(module, stack, buckets, size, priority) return _steal_pin(module, stack, buckets, size, priority)
module._pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)[offset:offset + size] module._pin = pin
module._pin.untyped_storage()._comfy_hostbuf = hostbuf
stack.append((module, offset)) stack.append((module, offset))
module._pin_registered = True module._pin_registered = True
module._pin_stack_index = len(stack) - 1 module._pin_stack_index = len(stack) - 1

2
requirements.txt
View File

@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0
filelock filelock
av>=16.0.0 av>=16.0.0
comfy-kitchen==0.2.10 comfy-kitchen==0.2.10
comfy-aimdo==0.4.8 comfy-aimdo==0.4.9
requests requests
simpleeval>=1.0.0 simpleeval>=1.0.0
blake3 blake3