11import torch
2+ import warnings
23import weakref
34
45import comfy_aimdo .model_vbar
@@ -28,6 +29,18 @@ def cleanup_prefetched_modules(module, comfy_modules):
2829 comfy_aimdo .model_vbar .vbar_unpin (module ._v_block )
2930 del module ._v_block_faulted
3031
32+ def _drop_graph (module ):
33+ graph = getattr (module , "_comfy_graph" , None )
34+ if graph is None :
35+ return
36+ # reset() through the bound method surfaces the allocator's benign
37+ # "uncaptured free of a captured allocation" as catchable Python warnings;
38+ # a plain del frees from the C++ dealloc path and spams stderr instead
39+ with warnings .catch_warnings ():
40+ warnings .simplefilter ("ignore" )
41+ graph ["graph" ].reset ()
42+ del module ._comfy_graph
43+
3144def cleanup_prefetch_queues ():
3245 global PREFETCH_QUEUES , GRAPH_CAPTURE_STREAMS
3346
@@ -41,7 +54,7 @@ def cleanup_prefetch_queues():
4154 cleanup_prefetched_modules (prefetched_module , comfy_modules )
4255 PREFETCH_QUEUES = []
4356 for module in GRAPH_MODULES :
44- del module . _comfy_graph
57+ _drop_graph ( module )
4558 GRAPH_MODULES .clear ()
4659 GRAPH_WARMED_MODULES .clear ()
4760 GRAPH_CAPTURE_STREAMS = {}
@@ -117,6 +130,7 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap
117130 if signature is not None :
118131 module ._v_block_faulted = True
119132 if signature is not None :
133+ _drop_graph (module )
120134 graph = torch .cuda .CUDAGraph ()
121135 if generator is not None :
122136 graph .register_generator_state (generator )
0 commit comments