xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-01 19:40:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge e148a7f1bb6ecd9f756526da0a93ac45dc78234c into b2ef58e2b17e73ca8cd376a1cdc976518ebbc168

Browse Source
This commit is contained in:
rattus 2025-11-24 11:08:55 -08:00 committed by GitHub
commit 6eb43de4b1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
comfy/model_management.py
View File

@ -1098,6 +1098,15 @@ if not args.disable_pinned_memory:
MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.95 MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.95
logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024))) logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))
def discard_cuda_async_error():
try:
a = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
b = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
_ = a + b
torch.cuda.synchronize()
except torch.AcceleratorError:
#Dump it! We already know about it from the synchronous return
pass
def pin_memory(tensor): def pin_memory(tensor):
global TOTAL_PINNED_MEMORY global TOTAL_PINNED_MEMORY
@ -1128,6 +1137,8 @@ def pin_memory(tensor):
PINNED_MEMORY[ptr] = size PINNED_MEMORY[ptr] = size
TOTAL_PINNED_MEMORY += size TOTAL_PINNED_MEMORY += size
return True return True
else:
discard_cuda_async_error()
return False return False
@ -1156,6 +1167,8 @@ def unpin_memory(tensor):
if len(PINNED_MEMORY) == 0: if len(PINNED_MEMORY) == 0:
TOTAL_PINNED_MEMORY = 0 TOTAL_PINNED_MEMORY = 0
return True return True
else:
discard_cuda_async_error()
return False return False