wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-10 21:42:37 +08:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

Try to disable pinned memory if wrong behavior detected.

Browse Source 
Not sure if this is going to improve anything
This commit is contained in:
comfyanonymous 2025-11-10 14:59:07 -05:00 committed by GitHub
parent c350009236
commit aa0a025df9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
comfy/model_management.py
View File

@ -1128,6 +1128,16 @@ def pin_memory(tensor):
return False
def disable_pinned():
global MAX_PINNED_MEMORY
global TOTAL_PINNED_MEMORY
MAX_PINNED_MEMORY = -1
logging.warning("Bad behavior detected, disabling pinned memory")
for ptr in PINNED_MEMORY: # TODO: check if this makes things worse
torch.cuda.cudart().cudaHostUnregister(ptr)
PINNED_MEMORY.clear()
TOTAL_PINNED_MEMORY = 0
def unpin_memory(tensor):
global TOTAL_PINNED_MEMORY
if MAX_PINNED_MEMORY <= 0:
@ -1142,6 +1152,7 @@ def unpin_memory(tensor):
size_stored = PINNED_MEMORY.get(ptr, None)
if size_stored is None:
logging.warning("Tried to unpin tensor not pinned by ComfyUI")
disable_pinned()
return False
if size != size_stored: