zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-24 18:10:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Implement Minimax Music 3 + Core Support for Cuda Graphs (#15570)

Browse Source
This commit is contained in:
rattus
2026-08-14 02:10:08 +10:00
committed by GitHub
parent 12666983cb
commit efd4e951a0
19 changed files with 1333 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

96
comfy/model_prefetch.py
View File

@@ -1,11 +1,18 @@
import torch
import weakref
import comfy_aimdo.model_vbar
from comfy.cli_args import args
import comfy.memory_management
import comfy.model_management
import comfy.ops
PREFETCH_QUEUES = []
GRAPH_MODULES = weakref.WeakSet()
GRAPH_WARMED_MODULES = weakref.WeakSet()
GRAPH_CAPTURE_STREAMS = {}
def cleanup_prefetched_modules(comfy_modules):
def cleanup_prefetched_modules(module, comfy_modules):
for s in comfy_modules:
prefetch = getattr(s, "_prefetch", None)
if prefetch is None:
@@ -17,39 +24,74 @@ def cleanup_prefetched_modules(comfy_modules):
if prefetch["signature"] is not None:
comfy_aimdo.model_vbar.vbar_unpin(s._v)
delattr(s, "_prefetch")
if getattr(module, "_v_block_faulted", False):
comfy_aimdo.model_vbar.vbar_unpin(module._v_block)
del module._v_block_faulted
def cleanup_prefetch_queues():
global PREFETCH_QUEUES
global PREFETCH_QUEUES, GRAPH_CAPTURE_STREAMS
for queue in PREFETCH_QUEUES:
for entry in queue:
if entry is None or not isinstance(entry, tuple):
continue
_, prefetch_state = entry
comfy_modules = prefetch_state[1]
prefetched_module, comfy_modules = prefetch_state
if comfy_modules is not None:
cleanup_prefetched_modules(comfy_modules)
cleanup_prefetched_modules(prefetched_module, comfy_modules)
PREFETCH_QUEUES = []
for module in GRAPH_MODULES:
del module._comfy_graph
GRAPH_MODULES.clear()
GRAPH_WARMED_MODULES.clear()
GRAPH_CAPTURE_STREAMS = {}
def prefetch_queue_pop(queue, device, module):
def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_graph=False, generator=None):
enable_graph = enable_graph and not args.disable_cuda_graphs and comfy.model_management.is_device_cuda(device)
if queue is None:
if core is not None:
core()
return
capture_stream = None
if enable_graph:
capture_stream = GRAPH_CAPTURE_STREAMS.get(device)
if capture_stream is None:
capture_stream = torch.cuda.Stream(device=device)
GRAPH_CAPTURE_STREAMS[device] = capture_stream
signature = None
graph_hit = False
graph = getattr(module, "_comfy_graph", None) if enable_graph else None
if graph is not None:
signature = comfy_aimdo.model_vbar.vbar_fault(module._v_block)
if signature is not None:
module._v_block_faulted = True
graph_hit = comfy_aimdo.model_vbar.vbar_signature_compare(signature, graph["signature"])
consumed = queue.pop(0)
if consumed is not None:
offload_stream, prefetch_state = consumed
if offload_stream is not None:
offload_stream.wait_stream(comfy.model_management.current_stream(device))
_, comfy_modules = prefetch_state
prefetched_module, comfy_modules = prefetch_state
if comfy_modules is not None:
cleanup_prefetched_modules(comfy_modules)
cleanup_prefetched_modules(prefetched_module, comfy_modules)
if graph_hit:
queue[0] = (None, (module, []))
graph["graph"].replay()
return
fully_faulted = False
prefetch = queue[0]
if prefetch is not None:
comfy_modules = []
for s in prefetch.modules():
if hasattr(s, "_v"):
comfy_modules.append(s)
prefetch_modules = prefetch if isinstance(prefetch, (list, tuple)) else (prefetch,)
for root in prefetch_modules:
for s in root.modules():
if hasattr(s, "_v"):
comfy_modules.append(s)
registerable_size = 0
for s in comfy_modules:
@@ -59,11 +101,41 @@ def prefetch_queue_pop(queue, device, module):
if lowvram_fn is not None:
registerable_size += lowvram_fn.memory_required()
offload_stream = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True)
offload_stream, fully_faulted = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True, return_faulted=True)
if not comfy.model_management.args.fast_disk:
comfy.model_management.ensure_pin_registerable(registerable_size)
comfy.model_management.sync_stream(device, offload_stream)
queue[0] = (offload_stream, (prefetch, comfy_modules))
if fully_faulted and dtype is not None:
for comfy_module in comfy_modules:
comfy.ops.resolve_cast_module_with_vbar(comfy_module, dtype, device, dtype, None, False, return_weights=False)
queue[0] = (offload_stream, (module, comfy_modules))
if core is not None:
if enable_graph and fully_faulted and module in GRAPH_WARMED_MODULES:
if signature is None:
signature = comfy_aimdo.model_vbar.vbar_fault(module._v_block)
if signature is not None:
module._v_block_faulted = True
if signature is not None:
graph = torch.cuda.CUDAGraph()
if generator is not None:
graph.register_generator_state(generator)
capture_stream.wait_stream(comfy.model_management.current_stream(device))
with torch.cuda.graph(graph, stream=capture_stream, capture_error_mode="thread_local"):
core()
comfy.model_management.current_stream(device).wait_stream(capture_stream)
graph.replay()
module._comfy_graph = {"graph": graph, "signature": signature}
GRAPH_MODULES.add(module)
return
if capture_stream is None:
core()
else:
capture_stream.wait_stream(comfy.model_management.current_stream(device))
with torch.cuda.stream(capture_stream):
core()
comfy.model_management.current_stream(device).wait_stream(capture_stream)
GRAPH_WARMED_MODULES.add(module)
def make_prefetch_queue(queue, device, transformer_options):
if (not transformer_options.get("prefetch_dynamic_vbars", False)