From 00d34d92fe0afbfbab3893ebbab2d5d70f5e9882 Mon Sep 17 00:00:00 2001 From: rattus <46076784+rattus128@users.noreply.github.com> Date: Wed, 9 Sep 2026 02:32:43 +1000 Subject: [PATCH] Comfy Aimdo 0.5.3 + Memory compiler fixes (#16180) --- comfy/latent_formats.py | 2 + comfy/ldm/lightricks/av_model.py | 2 +- comfy/ldm/minimax/model.py | 2 +- comfy/ldm/minimax_music/ar.py | 2 +- comfy/model_management.py | 4 - comfy/model_prefetch.py | 81 ++++++++++++------- comfy/multigpu.py | 28 ++++--- comfy/sd.py | 10 ++- comfy/text_encoders/llama.py | 2 +- comfy_extras/nodes_sparse_attention.py | 7 +- latent_preview.py | 13 ++- requirements.txt | 2 +- .../execution_test/preview_compiler_test.py | 57 +++++++++++++ 13 files changed, 153 insertions(+), 59 deletions(-) create mode 100644 tests-unit/execution_test/preview_compiler_test.py diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 6a60a7630..958bdacad 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -9,6 +9,7 @@ class LatentFormat: latent_rgb_factors_bias = None latent_rgb_factors_reshape = None taesd_decoder_name = None + compile_preview = False spacial_downscale_ratio = 8 temporal_downscale_ratio = 1 @@ -625,6 +626,7 @@ class MiniMaxH3Video(LatentFormat): temporal_downscale_ratio = 4 scale_factor = 1.0 taesd_decoder_name = "taeh3" + compile_preview = True latent_rgb_factors = [ [-0.018555, 0.024344, -0.017536], diff --git a/comfy/ldm/lightricks/av_model.py b/comfy/ldm/lightricks/av_model.py index d253b0144..baffc9ed2 100644 --- a/comfy/ldm/lightricks/av_model.py +++ b/comfy/ldm/lightricks/av_model.py @@ -938,7 +938,7 @@ class LTXAVModel(LTXVModel): stg_self_attn_blocks = transformer_options.get("stg_self_attn_blocks", ()) # Process transformer blocks - comfy.model_prefetch.malloc_graph_begin(self, vx.device) + comfy.model_prefetch.malloc_graph_begin(vx.device) for i, block in enumerate(self.transformer_blocks): comfy.model_prefetch.prefetch_queue_pop( prefetch_queue, vx.device, block, malloc_scope="block" diff --git a/comfy/ldm/minimax/model.py b/comfy/ldm/minimax/model.py index d7959027e..780df2435 100644 --- a/comfy/ldm/minimax/model.py +++ b/comfy/ldm/minimax/model.py @@ -573,7 +573,7 @@ class MiniMaxH3Model(nn.Module): compile_allocations = comfy.model_prefetch.malloc_graph_enabled(x[0].device) if compile_allocations: out = [torch.empty_like(x[0]), torch.empty_like(x[1])] - comfy.model_prefetch.malloc_graph_begin(self, x[0].device) + comfy.model_prefetch.malloc_graph_begin(x[0].device) graph_out = comfy.patcher_extension.WrapperExecutor.new_class_executor( self._forward, self, diff --git a/comfy/ldm/minimax_music/ar.py b/comfy/ldm/minimax_music/ar.py index 78a4c7c86..30b556a49 100644 --- a/comfy/ldm/minimax_music/ar.py +++ b/comfy/ldm/minimax_music/ar.py @@ -297,7 +297,7 @@ class MiniMaxMusic3AR(nn.Module): break if frame_index: - comfy.model_prefetch.malloc_graph_begin(self, device) + comfy.model_prefetch.malloc_graph_begin(device) c0, code_or_stop, stop_token = self._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask) if pending_code is None: pending_code = torch.empty_like(code_or_stop, device="cpu", pin_memory=cuda_device) diff --git a/comfy/model_management.py b/comfy/model_management.py index dd50f4c4e..e62fd3d74 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1385,7 +1385,6 @@ LARGEST_CASTED_WEIGHT = (None, 0) STREAM_AIMDO_CAST_BUFFERS = {} LARGEST_AIMDO_CASTED_WEIGHT = (None, 0) CROSS_STEP_STATE = weakref.WeakSet() -MALLOC_GRAPH_MODULES = weakref.WeakSet() DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3 @@ -1471,9 +1470,6 @@ def reset_cast_buffers(): STREAM_CAST_BUFFERS.clear() STREAM_AIMDO_CAST_BUFFERS.clear() - for module in MALLOC_GRAPH_MODULES: - del module._comfy_malloc_graph - MALLOC_GRAPH_MODULES.clear() soft_empty_cache() def get_offload_stream(device): diff --git a/comfy/model_prefetch.py b/comfy/model_prefetch.py index 35be626ae..d5952c8e9 100644 --- a/comfy/model_prefetch.py +++ b/comfy/model_prefetch.py @@ -1,9 +1,9 @@ -import contextlib import logging import threading import warnings import weakref +import comfy_kitchen as ck import torch import comfy_aimdo.malloc_graph @@ -14,11 +14,11 @@ import comfy.model_management import comfy.ops PREFETCH_QUEUES = [] -GRAPH_MODULES = weakref.WeakSet() GRAPH_WARMED_MODULES = weakref.WeakSet() GRAPH_CAPTURE_STREAMS = {} -ACTIVE_MALLOC_GRAPHS = {} +MALLOC_GRAPHS = {} MALLOC_GRAPH_BREAKS = 0 +MALLOC_GRAPH_ROGUES = 0 MALLOC_GRAPH_USED = False def _malloc_graph_break(): @@ -29,40 +29,62 @@ def _malloc_graph_break(): def malloc_graph_enabled(device): return not args.disable_comfy_compiler and comfy.memory_management.aimdo_enabled and comfy.model_management.is_device_cuda(device) -@contextlib.contextmanager -def pause_malloc_graph(sync=False): - graph = ACTIVE_MALLOC_GRAPHS.get(threading.get_ident()) - if graph is not None: - graph.pause(sync=sync) - try: - yield - finally: - if graph is not None: - graph.resume(sync=sync) +class _PauseMallocGraph: + def __init__(self, sync=False): + self.sync = sync -def malloc_graph_begin(module, device): + def __enter__(self): + graph = MALLOC_GRAPHS.get(threading.get_ident()) + if graph is not None and graph._comfy_active: + graph.pause(sync=self.sync) + + def __exit__(self, *args): + graph = MALLOC_GRAPHS.get(threading.get_ident()) + if graph is not None and graph._comfy_active: + graph.resume(sync=self.sync) + +def pause_malloc_graph(sync=False): + return _PauseMallocGraph(sync) + +def malloc_graph_begin(device): global MALLOC_GRAPH_USED if not malloc_graph_enabled(device): return - graph = getattr(module, "_comfy_malloc_graph", None) + thread_id = threading.get_ident() + graph = MALLOC_GRAPHS.get(thread_id) if graph is None: graph = comfy_aimdo.malloc_graph.record( comfy.model_management.current_stream(device), args.assert_graph_breaks ) - module._comfy_malloc_graph = graph - comfy.model_management.MALLOC_GRAPH_MODULES.add(module) + graph._comfy_cuda_graph_modules = weakref.WeakSet() + MALLOC_GRAPHS[thread_id] = graph else: graph.push() - ACTIVE_MALLOC_GRAPHS[threading.get_ident()] = graph + if hasattr(ck, "set_allocation_context"): + ck.set_allocation_context(pause_malloc_graph()) + graph._comfy_active = True MALLOC_GRAPH_USED = True def malloc_graph_end(): thread_id = threading.get_ident() - graph = ACTIVE_MALLOC_GRAPHS.get(thread_id) - if graph is not None: + graph = MALLOC_GRAPHS.get(thread_id) + if graph is not None and graph._comfy_active: if graph.pop(): _malloc_graph_break() - ACTIVE_MALLOC_GRAPHS.pop(thread_id) + graph._comfy_active = False + +def cleanup_malloc_graph(): + global MALLOC_GRAPH_ROGUES + + graph = MALLOC_GRAPHS.pop(threading.get_ident(), None) + if graph is not None: + if graph._comfy_active: + graph.abort() + graph._comfy_active = False + for module in graph._comfy_cuda_graph_modules: + _drop_graph(module) + MALLOC_GRAPH_ROGUES += graph.rogue_count + del graph def cleanup_prefetched_modules(module, comfy_modules): for s in comfy_modules: @@ -95,11 +117,10 @@ def _drop_graph(module): def cleanup_prefetch_queues(): global PREFETCH_QUEUES global MALLOC_GRAPH_BREAKS + global MALLOC_GRAPH_ROGUES global MALLOC_GRAPH_USED - graph = ACTIVE_MALLOC_GRAPHS.pop(threading.get_ident(), None) - if graph is not None: - graph.abort() + cleanup_malloc_graph() for queue in PREFETCH_QUEUES: for entry in queue: if entry is None or not isinstance(entry, tuple): @@ -109,17 +130,17 @@ def cleanup_prefetch_queues(): if comfy_modules is not None: cleanup_prefetched_modules(prefetched_module, comfy_modules) PREFETCH_QUEUES = [] - for module in GRAPH_MODULES: - _drop_graph(module) - GRAPH_MODULES.clear() GRAPH_WARMED_MODULES.clear() if MALLOC_GRAPH_USED: - logging.info("Comfy model compiler graph breaks: %d", MALLOC_GRAPH_BREAKS) + logging.info("Comfy model compiler graph breaks: %d, rogues: %d", MALLOC_GRAPH_BREAKS, MALLOC_GRAPH_ROGUES) MALLOC_GRAPH_BREAKS = 0 + MALLOC_GRAPH_ROGUES = 0 MALLOC_GRAPH_USED = False def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_graph=False, generator=None, malloc_scope=None): - malloc_graph = ACTIVE_MALLOC_GRAPHS.get(threading.get_ident()) + malloc_graph = MALLOC_GRAPHS.get(threading.get_ident()) + if malloc_graph is not None and not malloc_graph._comfy_active: + malloc_graph = None enable_graph = enable_graph and malloc_graph is not None and not args.disable_cuda_graphs and comfy.model_management.is_device_cuda(device) and getattr(module, "_v_block", None) is not None if queue is None: if malloc_graph is not None and malloc_scope is not None: @@ -223,7 +244,7 @@ def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_grap comfy.model_management.current_stream(device).wait_stream(capture_stream) graph.replay() module._comfy_graph = {"graph": graph, "signature": signature} - GRAPH_MODULES.add(module) + malloc_graph._comfy_cuda_graph_modules.add(module) return if capture_stream is None: core() diff --git a/comfy/multigpu.py b/comfy/multigpu.py index 2b6d8260d..57644af71 100644 --- a/comfy/multigpu.py +++ b/comfy/multigpu.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: import comfy.utils import comfy.patcher_extension import comfy.model_management +import comfy.model_prefetch class MultiGPUThreadPool: @@ -46,18 +47,21 @@ class MultiGPUThreadPool: return result_q.put((None, e)) return - while True: - item = work_q.get() - if item is None: - break - fn, args, kwargs = item - try: - result = fn(*args, **kwargs) - result_q.put((result, None)) - except comfy.model_management.InterruptProcessingException as e: - result_q.put((None, e)) - except Exception as e: - result_q.put((None, e)) + try: + while True: + item = work_q.get() + if item is None: + break + fn, args, kwargs = item + try: + result = fn(*args, **kwargs) + result_q.put((result, None)) + except comfy.model_management.InterruptProcessingException as e: + result_q.put((None, e)) + except Exception as e: + result_q.put((None, e)) + finally: + comfy.model_prefetch.cleanup_malloc_graph() def submit(self, device: torch.device, fn, *args, **kwargs): self._work_queues[device].put((fn, args, kwargs)) diff --git a/comfy/sd.py b/comfy/sd.py index a73607cb0..54fbbd89c 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -35,6 +35,7 @@ import os import comfy.utils import comfy.ops +import comfy.model_prefetch from . import clip_vision from . import gligen @@ -1227,7 +1228,8 @@ class VAE: with model_management.cuda_device_context(self.device): try: memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype) - model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload) + with comfy.model_prefetch.pause_malloc_graph(): + model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload) free_memory = self.patcher.get_free_memory(self.device) batch_number = int(free_memory / memory_used) batch_number = max(1, batch_number) @@ -1235,7 +1237,8 @@ class VAE: # Pre-allocate output for VAEs that support direct buffer writes preallocated = False if getattr(self.first_stage_model, 'comfy_has_chunked_io', False): - pixel_samples = torch.empty(self.first_stage_model.decode_output_shape(samples_in.shape), device=self.output_device, dtype=self.vae_output_dtype()) + with comfy.model_prefetch.pause_malloc_graph(): + pixel_samples = torch.empty(self.first_stage_model.decode_output_shape(samples_in.shape), device=self.output_device, dtype=self.vae_output_dtype()) preallocated = True for x in range(0, samples_in.shape[0], batch_number): @@ -1245,7 +1248,8 @@ class VAE: else: out = self.first_stage_model.decode(samples, **vae_options).to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True) if pixel_samples is None: - pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device, dtype=self.vae_output_dtype()) + with comfy.model_prefetch.pause_malloc_graph(): + pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device, dtype=self.vae_output_dtype()) pixel_samples[x:x+batch_number].copy_(out) del out self.process_output(pixel_samples[x:x+batch_number]) diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index a61c5adc7..c7147904b 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -1036,7 +1036,7 @@ class BaseGenerate: for step in tqdm(range(max_length), desc="Generating tokens"): if step > 0: if compile_allocations: - comfy.model_prefetch.malloc_graph_begin(self, device) + comfy.model_prefetch.malloc_graph_begin(device) embeds = self.model.embed_tokens(decode_tokens).to(execution_dtype) current_input_ids = decode_tokens if initial_input_ids is not None else None position_ids = torch.tensor([[next_pos]], device=device) if next_pos is not None else None diff --git a/comfy_extras/nodes_sparse_attention.py b/comfy_extras/nodes_sparse_attention.py index 441b474c9..006d1eb35 100644 --- a/comfy_extras/nodes_sparse_attention.py +++ b/comfy_extras/nodes_sparse_attention.py @@ -5,6 +5,7 @@ from __future__ import annotations import logging import re +import weakref import comfy_kitchen as ck import torch @@ -142,11 +143,11 @@ class SparseAttnPatch: def vsa_rope_freqs(self, rope_freqs, plan): hit = self.vsa_rope - if hit is not None and hit[0] is rope_freqs and hit[1] is plan: + if hit is not None and hit[0]() is rope_freqs and hit[1] is plan: return hit[2] padded = rope_freqs.new_zeros((1, plan["n"]) + tuple(rope_freqs.shape[2:])) padded[0, plan["inv"]] = rope_freqs[0] - self.vsa_rope = (rope_freqs, plan, padded) + self.vsa_rope = (weakref.ref(rope_freqs), plan, padded) return padded @@ -257,6 +258,7 @@ def h3_sparse_attention(attn, x, rope_freqs, transformer_options, patch: SparseA if patch.vsa: plan = patch.vsa_plan(transformer_options["minimax_h3_layout"], x.device) n = plan["n"] + freqs = patch.vsa_rope_freqs(rope_freqs, plan) key = (block_index, n, tuple(transformer_options.get("uuids", ()))) # statistics per conditioning branch pooled = patch.pooled.get(key) @@ -268,7 +270,6 @@ def h3_sparse_attention(attn, x, rope_freqs, transformer_options, patch: SparseA ) if patch.vsa: - freqs = patch.vsa_rope_freqs(rope_freqs, plan) sink = sink_q = (0, plan["n_prefix"]) extra = {"tail": False, "block_len": plan["block_len"]} gate = attn.to_gate_compress diff --git a/latent_preview.py b/latent_preview.py index d98b70019..e3864e938 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -4,6 +4,7 @@ from comfy.cli_args import args, LatentPreviewMethod from comfy.taesd.taesd import TAESD from comfy.sd import VAE import comfy.model_management +import comfy.model_prefetch import folder_paths import comfy.utils import logging @@ -45,8 +46,16 @@ class TAESDPreviewerImpl(LatentPreviewer): return preview_to_image(x_sample) class TAEHVPreviewerImpl(TAESDPreviewerImpl): + def __init__(self, taesd, compile_preview=False): + super().__init__(taesd) + self.compile_preview = compile_preview + def decode_latent_to_preview(self, x0): - x_sample = self.taesd.decode(x0[:1, :, :1])[0][0] + samples = x0[:1, :, :1] + if self.compile_preview and comfy.model_prefetch.malloc_graph_enabled(self.taesd.device): + comfy.model_prefetch.malloc_graph_begin(self.taesd.device) + x_sample = self.taesd.decode(samples)[0][0] + comfy.model_prefetch.malloc_graph_end() return preview_to_image(x_sample, do_scale=False) class Latent2RGBPreviewer(LatentPreviewer): @@ -97,7 +106,7 @@ def get_previewer(device, latent_format): if latent_format.taesd_decoder_name in VIDEO_TAES: taesd = VAE(comfy.utils.load_torch_file(taesd_decoder_path)) taesd.first_stage_model.show_progress_bar = False - previewer = TAEHVPreviewerImpl(taesd) + previewer = TAEHVPreviewerImpl(taesd, compile_preview=latent_format.compile_preview) else: taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device) previewer = TAESDPreviewerImpl(taesd) diff --git a/requirements.txt b/requirements.txt index 8a12a2266..be4958429 100644 --- a/requirements.txt +++ b/requirements.txt @@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0 filelock av>=17.0.0 comfy-kitchen==0.2.33 -comfy-aimdo==0.5.2 +comfy-aimdo==0.5.3 requests simpleeval>=1.0.0 blake3 diff --git a/tests-unit/execution_test/preview_compiler_test.py b/tests-unit/execution_test/preview_compiler_test.py new file mode 100644 index 000000000..dc9664ad3 --- /dev/null +++ b/tests-unit/execution_test/preview_compiler_test.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock, Mock + +import pytest + +import latent_preview +from comfy import latent_formats + + +def test_minimax_h3_enables_preview_compiler(): + assert latent_formats.MiniMaxH3Video.compile_preview + assert latent_formats.MiniMaxH3AV.compile_preview + assert not latent_formats.HunyuanVideo.compile_preview + + +def test_video_preview_compiles_decode(monkeypatch): + taesd = Mock() + taesd.device = "cuda:0" + taesd.decode.return_value = [[Mock()]] + previewer = latent_preview.TAEHVPreviewerImpl(taesd, compile_preview=True) + x0 = MagicMock() + samples = Mock(shape=(1, 24, 1, 30, 52)) + x0.__getitem__.return_value = samples + + monkeypatch.setattr(latent_preview, "preview_to_image", Mock()) + monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_enabled", Mock(return_value=True)) + calls = [] + taesd.decode.side_effect = lambda value: calls.append(("decode", value)) or [[Mock()]] + begin = Mock(side_effect=lambda device: calls.append(("begin", device))) + end = Mock(side_effect=lambda: calls.append(("end",))) + monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_begin", begin) + monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_end", end) + + previewer.decode_latent_to_preview(x0) + assert calls == [ + ("begin", "cuda:0"), + ("decode", samples), + ("end",), + ] + + +def test_video_preview_leaves_failed_compiler_scope_for_execution_cleanup(monkeypatch): + taesd = Mock() + taesd.device = "cuda:0" + taesd.decode.side_effect = RuntimeError("decode failed") + previewer = latent_preview.TAEHVPreviewerImpl(taesd, compile_preview=True) + x0 = MagicMock() + + monkeypatch.setattr(latent_preview, "preview_to_image", Mock()) + monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_enabled", Mock(return_value=True)) + monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_begin", Mock()) + end = Mock() + monkeypatch.setattr(latent_preview.comfy.model_prefetch, "malloc_graph_end", end) + + with pytest.raises(RuntimeError, match="decode failed"): + previewer.decode_latent_to_preview(x0) + + end.assert_not_called()