wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-22 23:17:34 +08:00
Code Issues Actions 40 Packages Projects Releases Wiki Activity

Rabbit

Browse Source
This commit is contained in:
Rattus 2026-04-30 09:14:14 +10:00
parent dfdb4071a6
commit f3ad2b7f2e
1 changed files with 2 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
comfy/model_prefetch.py
View File

@ -4,7 +4,6 @@ import comfy.ops
PREFETCH_QUEUES = [] PREFETCH_QUEUES = []
def cleanup_prefetched_modules(comfy_modules): def cleanup_prefetched_modules(comfy_modules):
for s in comfy_modules: for s in comfy_modules:
prefetch = getattr(s, "_prefetch", None) prefetch = getattr(s, "_prefetch", None)
@ -17,12 +16,13 @@ def cleanup_prefetched_modules(comfy_modules):
if prefetch["signature"] is not None: if prefetch["signature"] is not None:
comfy_aimdo.model_vbar.vbar_unpin(s._v) comfy_aimdo.model_vbar.vbar_unpin(s._v)
delattr(s, "_prefetch") delattr(s, "_prefetch")
def cleanup_prefetch_queues(): def cleanup_prefetch_queues():
global PREFETCH_QUEUES global PREFETCH_QUEUES
for queue in PREFETCH_QUEUES: for queue in PREFETCH_QUEUES:
for entry in queue: for entry in queue:
if entry is None: if entry is None or not isinstance(entry, tuple):
continue continue
_, prefetch_state = entry _, prefetch_state = entry
comfy_modules = prefetch_state[1] comfy_modules = prefetch_state[1]
@ -30,7 +30,6 @@ def cleanup_prefetch_queues():
cleanup_prefetched_modules(comfy_modules) cleanup_prefetched_modules(comfy_modules)
PREFETCH_QUEUES = [] PREFETCH_QUEUES = []
def prefetch_queue_pop(queue, device, module): def prefetch_queue_pop(queue, device, module):
if queue is None: if queue is None:
return return
@ -51,9 +50,6 @@ def prefetch_queue_pop(queue, device, module):
comfy_modules.append(s) comfy_modules.append(s)
offload_stream = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True) offload_stream = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True)
if offload_stream is None:
queue[0] = None
return
comfy.model_management.sync_stream(device, offload_stream) comfy.model_management.sync_stream(device, offload_stream)
queue[0] = (offload_stream, (prefetch, comfy_modules)) queue[0] = (offload_stream, (prefetch, comfy_modules))