Skip to content

Commit ad2d920

Browse files
authored
Merge branch 'master' into fix/ltx2-prompt-empty-output
2 parents 962a262 + 7fe8a61 commit ad2d920

4 files changed

Lines changed: 261 additions & 22 deletions

File tree

comfy/model_prefetch.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import torch
2+
import warnings
23
import weakref
34

45
import 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+
3144
def 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)

comfy/ops.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1639,7 +1639,7 @@ def forward_comfy_cast_weights(self, input, out_dtype=None):
16391639
self.norm_type, self.scale_grad_by_freq, self.sparse)
16401640
target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
16411641
x = x.to(dtype=target_dtype)
1642-
if scale is not None and scale != 1.0:
1642+
if scale is not None:
16431643
x = x * scale.to(dtype=target_dtype)
16441644
return x
16451645

0 commit comments

Comments
 (0)