Leak report: say when ComfyUI's parent hand-over finds no parent (and whether gc was running), and which patchers still hold the model

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

core/leak_report.py

@@ -55,16 +55,59 @@ def report(mm, when) -> int:
         model = m.real_model()
         if model is None:
             continue
-        log.warning("FunUI Leak", f"{when}: {type(model).__name__} outlived its patcher. What holds it:\n" + "\n".join(chains(model)))
+        held = [r for r in gc.get_referrers(model) if type(r).__name__ == "dict" and "patches" in r]
+        owners = [o for h in held for o in gc.get_referrers(h) if type(o).__name__.startswith("ModelPatcher")]
+
+        def depth(p):
+            n = 0
+            while getattr(p, "parent", None) is not None and n < 50:
+                p, n = p.parent, n + 1
+            return n
+        lineage = ", ".join(f"{type(o).__name__}(parents {depth(o)}, tracked {any(l.model is o for l in mm.current_loaded_models)})" for o in owners)
+        log.warning("FunUI Leak", f"{when}: {type(model).__name__} outlived its patcher. Patchers holding it: {lineage or 'none'}. What holds it:\n" + "\n".join(chains(model)))
     return len(dead)
 
 
+_in_gc = [False]
+
+
+def _gc_flag(phase, info):
+    _in_gc[0] = phase == "start"
+
+
+def watch_switch(mm) -> bool:
+    """ComfyUI follows a model through its patcher's parent when the patcher dies. Say when that hand-over finds no parent, and whether
+    Python's own gc was running at that moment (it clears weak links to every patcher in a cycle before the hand-over can use them)."""
+    original = getattr(mm.LoadedModel, "_switch_parent", None)
+    if original is None or getattr(original, TAG, False):
+        return False
+    if _gc_flag not in gc.callbacks:
+        gc.callbacks.append(_gc_flag)
+
+    def _switch_parent(self):
+        try:
+            parent = self._parent_model()
+            real = self.real_model() if callable(self.real_model) else None
+            if parent is None and real is not None:
+                live = [type(o).__name__ for o in gc.get_referrers(real) if type(o).__name__.startswith("ModelPatcher")]
+                log.warning("FunUI Leak", f"a {type(real).__name__} patcher died and its parent was already gone, so ComfyUI loses track of the model. "
+                            f"Python's gc was running: {_in_gc[0]}. Patchers still holding the model: {len(live)}")
+        except Exception as exc:            # noqa: BLE001 -- a diagnostic must never stop a run
+            log.failed("leak report", exc)
+        return original(self)
+
+    setattr(_switch_parent, TAG, True)
+    mm.LoadedModel._switch_parent = _switch_parent
+    return True
+
+
 def install(mm=None) -> bool:
     if mm is None:
         import comfy.model_management as mm
     original = getattr(mm, "cleanup_models_gc", None)
     if original is None or getattr(original, TAG, False):
         return False
+    watch_switch(mm)
 
     def cleanup_models_gc():
         try:

core/tests/test_leak_report.py

@@ -37,10 +37,40 @@ def test_a_model_kept_by_something_else_is_named_before_and_after_gc(comfyui, ca
 
 def test_install_wraps_once_and_a_clean_state_logs_nothing(comfyui, capsys):
     import comfy.model_management as mm
-    original = mm.cleanup_models_gc
+    original, switch = mm.cleanup_models_gc, mm.LoadedModel._switch_parent
     try:
         assert leak_report.install(mm) and not leak_report.install(mm)
         mm.cleanup_models_gc()
         assert "FunUI Leak" not in capsys.readouterr().err
     finally:
-        mm.cleanup_models_gc = original
+        mm.cleanup_models_gc, mm.LoadedModel._switch_parent = original, switch
+        if leak_report._gc_flag in gc.callbacks:
+            gc.callbacks.remove(leak_report._gc_flag)
+
+
+def test_a_hand_over_that_finds_no_parent_is_said(comfyui, capsys):
+    import weakref
+    import comfy.model_management as mm
+    import comfy.model_patcher as mp
+    original = mm.LoadedModel._switch_parent
+    try:
+        assert leak_report.watch_switch(mm) and not leak_report.watch_switch(mm)
+        base = torch.nn.Linear(2, 2)
+        base.model_config = type("C", (), {})()
+        parent = mp.ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
+        child = parent.clone()
+        loaded = mm.LoadedModel(child)
+        loaded.real_model = weakref.ref(base)
+        keep = Holder(base)
+        del parent
+        gc.collect()                    # the parent is held by the child: still there
+        assert loaded._parent_model() is not None
+        loaded._parent_model = lambda: None          # what a gc that cleared the weak link leaves behind
+        loaded._switch_parent()
+        shown = capsys.readouterr().err
+        assert "parent was already gone" in shown and "Python's gc was running: False" in shown
+        del keep
+    finally:
+        mm.LoadedModel._switch_parent = original
+        if leak_report._gc_flag in gc.callbacks:
+            gc.callbacks.remove(leak_report._gc_flag)