ComfyUI loses track of a model whose patcher died when the parent link is gone, though cached patchers still hold it: point such an entry at a live patcher before ComfyUI's leak check

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

core/leak_report.py

@@ -49,6 +49,40 @@ def chains(target, depth=DEPTH, nodes=NODES, paths=PATHS) -> list:
     return out
 
 
+def _owners(model):
+    """The live patchers whose model this is, deepest clone first."""
+    held = [r for r in gc.get_referrers(model) if type(r).__name__ == "dict" and "patches" in r]
+    found = {id(o): o for h in held for o in gc.get_referrers(h) if type(o).__name__.startswith("ModelPatcher") and getattr(o, "model", None) is model}
+
+    def depth(p):
+        n = 0
+        while getattr(p, "parent", None) is not None and n < 50:
+            p, n = p.parent, n + 1
+        return n
+    return sorted(found.values(), key=depth, reverse=True)
+
+
+def relink(mm) -> int:
+    """ComfyUI follows a loaded model through its patcher, and when that patcher dies it hands over to the patcher's direct parent
+    only. A parent that is gone leaves the entry dead although other patchers of the same model are alive (the loader's, the
+    LoRA's, the modifiers' outputs, all cached by ComfyUI): it then warns "memory leak" on every load. Such an entry is pointed
+    at the deepest live patcher of its model, so the model is tracked again. An entry with no live patcher stays dead and is
+    reported. -> how many were re-linked."""
+    n = 0
+    for m in list(mm.current_loaded_models):
+        if not m.is_dead():
+            continue
+        model = m.real_model()
+        owners = _owners(model) if model is not None else []
+        if not owners:
+            continue
+        m._set_model(owners[0])
+        m.device = owners[0].load_device
+        n += 1
+        log.info("FunUI Leak", f"{type(model).__name__} had lost its patcher; ComfyUI tracks it again through one of the {len(owners)} patchers still holding it")
+    return n
+
+
 def report(mm, when) -> int:
     dead = [m for m in mm.current_loaded_models if m.is_dead()]
     if dead:
@@ -122,6 +156,7 @@ def install(mm=None) -> bool:
 
     def cleanup_models_gc():
         try:
+            relink(mm)
             report(mm, "before ComfyUI's gc")
         except Exception as exc:            # noqa: BLE001 -- a diagnostic must never stop a run
             log.failed("leak report", exc)

core/tests/test_leak_report.py

@@ -74,3 +74,51 @@ def test_a_hand_over_that_finds_no_parent_is_said(comfyui, capsys):
         mm.LoadedModel._switch_parent = original
         if leak_report._gc_flag in gc.callbacks:
             gc.callbacks.remove(leak_report._gc_flag)
+
+
+def test_an_entry_whose_patcher_died_without_a_hand_over_is_tracked_again_through_a_live_one(comfyui, capsys):
+    """The state the rental showed: the cache keeps the loader, LoRA and modifier patchers (one chain), the sampling clone died and
+    ComfyUI could not hand over to its parent. The entry must be tracked again, and ComfyUI's own check must stay quiet."""
+    import weakref
+    import comfy.model_management as mm
+    import comfy.model_patcher as mp
+    base = torch.nn.Linear(2, 2)
+    base.model_config = type("C", (), {})()
+    root = mp.ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
+    chain = [root]
+    for _ in range(3):
+        chain.append(chain[-1].clone())
+    sampling = chain[-1].clone()
+    loaded = mm.LoadedModel(sampling)
+    loaded.real_model = weakref.ref(base)
+    mm.current_loaded_models.append(loaded)
+    try:
+        loaded._parent_model = lambda: None          # the hand-over finds no parent
+        del sampling
+        gc.collect()
+        assert loaded.is_dead()
+        assert leak_report.relink(mm) == 1 and not loaded.is_dead() and loaded.model is chain[-1]
+        assert "tracks it again" in capsys.readouterr().err
+        assert leak_report.relink(mm) == 0
+    finally:
+        mm.current_loaded_models.remove(loaded)
+
+
+def test_an_entry_with_no_live_patcher_is_left_for_the_report(comfyui):
+    import weakref
+    import comfy.model_management as mm
+    import comfy.model_patcher as mp
+    base = torch.nn.Linear(2, 2)
+    base.model_config = type("C", (), {})()
+    patcher = mp.ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
+    loaded = mm.LoadedModel(patcher)
+    loaded.real_model = weakref.ref(base)
+    keep = Holder(base)
+    mm.current_loaded_models.append(loaded)
+    try:
+        del patcher
+        gc.collect()
+        assert loaded.is_dead() and leak_report.relink(mm) == 0 and loaded.is_dead()
+        del keep
+    finally:
+        mm.current_loaded_models.remove(loaded)