fix(sampler): sweep leaked hooks below the blocks; log a hook census per run

v2a scale and bounded attention install on attn2 / video_to_audio_attn, which the
block-only sweep never reached and never tagged. The census prints what is still
installed before sampling, so a count that grows is proof rather than a hunch.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

ltx_enhancements.py

@@ -1262,16 +1262,30 @@ def _funpack_locate_blocks(model):
     return None
 
 
+def funpack_diffusion_model(model):
+    """The shared nn.Module every ModelPatcher clone points at."""
+    return getattr(getattr(model, "model", None), "diffusion_model", None)
+
+
+def _hook_stores(module):
+    return (getattr(module, "_forward_hooks", None),
+            getattr(module, "_forward_pre_hooks", None))
+
+
 def strip_funpack_block_hooks(model):
-    """Remove forward / forward-pre hooks this module installed on the shared
-    diffusion blocks in a previous run. Idempotent; safe to call every run."""
-    bl = _funpack_locate_blocks(model)
-    if bl is None:
+    """Remove forward / forward-pre hooks FunPack installed on the shared diffusion model
+    in a previous run. Idempotent; safe to call every run.
+
+    Sweeps every submodule, not only the blocks: the v2a scale sits on each block's
+    `video_to_audio_attn` and bounded attention on its `attn2`, so a block-only sweep left
+    those to accumulate silently. Foreign hooks are never touched — only tagged ones.
+    """
+    diff = funpack_diffusion_model(model)
+    if diff is None:
         return 0
     removed = 0
-    for block in bl:
-        for store in (getattr(block, "_forward_hooks", None),
-                      getattr(block, "_forward_pre_hooks", None)):
+    for module in diff.modules():
+        for store in _hook_stores(module):
             if not store:
                 continue
             for hid in [hid for hid, fn in list(store.items())
@@ -1279,10 +1293,29 @@ def strip_funpack_block_hooks(model):
                 store.pop(hid, None)
                 removed += 1
     if removed:
-        print(f"[FunPackEnhancements] Stripped {removed} leaked block hook(s) from a previous run")
+        print(f"[FunPackEnhancements] Stripped {removed} leaked hook(s) from a previous run")
     return removed
 
 
+def count_module_hooks(model):
+    """(hooks, modules) currently installed anywhere under the shared diffusion model.
+
+    Read at the start of a run, when FunPack has installed nothing yet, this should be 0.
+    A number that climbs generation after generation IS the progressive-corruption bug —
+    each leftover hook keeps rewriting activations with a previous run's state.
+    """
+    diff = funpack_diffusion_model(model)
+    if diff is None:
+        return 0, 0
+    hooks = modules = 0
+    for module in diff.modules():
+        n = sum(len(store) for store in _hook_stores(module) if store)
+        if n:
+            hooks += n
+            modules += 1
+    return hooks, modules
+
+
 # ---------------------------------------------------------------------------
 # Main enhancement builder
 # ---------------------------------------------------------------------------

samplers.py

@@ -1035,6 +1035,20 @@ def _tag_scene_wrapper(wrapper, prev):
     return wrapper
 
 
+def _tag_funpack_hook(fn):
+    """Mark a forward hook as FunPack's, so the run-start sweep can remove one that leaked.
+
+    Anything installed on the shared diffusion model needs this: the sweep only removes
+    what it can prove is ours, so an untagged hook is one nothing will ever clean up.
+    """
+    try:
+        from .ltx_enhancements import _FUNPACK_HOOK_TAG
+    except ImportError:
+        from ltx_enhancements import _FUNPACK_HOOK_TAG
+    setattr(fn, _FUNPACK_HOOK_TAG, True)
+    return fn
+
+
 def _strip_funpack_scene_wrappers(model):
     """Unwind FunPack per-scene wrappers leaked by a previous interrupted/failed run.
     Walks the recorded prev-wrapper chain back to the first non-FunPack wrapper (or
@@ -5623,6 +5637,7 @@ class FunPackLTXAVSceneChainSampler:
                 return (output[0] * s,) + tuple(output[1:])
             return output * s
 
+        _tag_funpack_hook(_hook)
         handles = []
         for blk in blocks:
             sub = getattr(blk, "video_to_audio_attn", None)
@@ -6113,6 +6128,7 @@ class FunPackLTXAVSceneChainSampler:
             except Exception:
                 return args, kwargs
 
+        _tag_funpack_hook(_hook)
         handles = []
         for blk in blocks:
             sub = getattr(blk, "attn2", None)
@@ -6298,10 +6314,19 @@ class FunPackLTXAVSceneChainSampler:
         # build_enhancements, so stale hooks can't fire on an unenhanced generation.
         try:
             try:
-                from .ltx_enhancements import strip_funpack_block_hooks
+                from .ltx_enhancements import strip_funpack_block_hooks, count_module_hooks
             except ImportError:
-                from ltx_enhancements import strip_funpack_block_hooks
+                from ltx_enhancements import strip_funpack_block_hooks, count_module_hooks
             strip_funpack_block_hooks(model)
+            # Nothing of ours is installed yet at this point, so whatever is left is either a
+            # third-party hook or a leak this sweep cannot prove is ours. Printing the count
+            # every run makes the difference visible: a number that climbs run after run is
+            # the progressive-degradation bug, not a hunch.
+            _left, _mods = count_module_hooks(model)
+            if _left:
+                print(f"[FunPackSceneChain] hook census before sampling: {_left} hook(s) on "
+                      f"{_mods} module(s) — FunPack installs none until it samples, so a count "
+                      f"that grows every run means hooks are leaking")
         except Exception as _e:
             print(f"[FunPackLTXAVSceneChainSampler] hook strip failed: {_e}")
         # Same defense for per-scene model_function_wrappers: normally unwound in the

tests/test_scene_wrapper_hygiene.py

@@ -142,3 +142,51 @@ def test_output_guidance_high_sigma_is_passthrough():
     out = wrap(apply_fn, {"input": denoised, "timestep": torch.tensor([0.9]),
                           "c": {}})
     assert torch.equal(out, denoised)
+
+
+# ── the sweep has to reach every module a hook can be installed on ────────────
+
+class _Block(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn2 = torch.nn.Identity()               # bounded attention hooks here
+        self.video_to_audio_attn = torch.nn.Identity() # v2a scale hooks here
+
+    def forward(self, x):
+        return x
+
+
+class _Diffusion(torch.nn.Module):
+    def __init__(self, n=2):
+        super().__init__()
+        self.transformer_blocks = torch.nn.ModuleList(_Block() for _ in range(n))
+
+    def forward(self, x):
+        return x
+
+
+def _patcher():
+    return types.SimpleNamespace(model=types.SimpleNamespace(diffusion_model=_Diffusion()))
+
+
+def test_the_sweep_reaches_hooks_installed_below_the_blocks():
+    """v2a sits on video_to_audio_attn and bounded attention on attn2 — a block-only sweep
+    left both to stack up run after run."""
+    import ltx_enhancements as enh
+    model = _patcher()
+    hook = samplers._tag_funpack_hook(lambda *a: None)
+    for blk in model.model.diffusion_model.transformer_blocks:
+        blk.attn2.register_forward_pre_hook(hook)
+        blk.video_to_audio_attn.register_forward_hook(hook)
+    assert enh.count_module_hooks(model) == (4, 4)
+    assert enh.strip_funpack_block_hooks(model) == 4
+    assert enh.count_module_hooks(model) == (0, 0)
+
+
+def test_the_sweep_leaves_hooks_that_are_not_funpacks():
+    import ltx_enhancements as enh
+    model = _patcher()
+    model.model.diffusion_model.transformer_blocks[0].attn2.register_forward_hook(
+        lambda *a: None)
+    assert enh.strip_funpack_block_hooks(model) == 0
+    assert enh.count_module_hooks(model) == (1, 1)