Comfy Aimdo 0.5.3 + Memory compiler fixes (#16180)

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],

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"

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,

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)

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):

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()

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))

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])

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

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

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)

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

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()