diff --git a/ltx_enhancements.py b/ltx_enhancements.py index fc5cda57..bdcd64d3 100644 --- a/ltx_enhancements.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/samplers.py b/samplers.py index 0a2fc272..a5f313fc 100644 --- a/samplers.py +++ b/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 diff --git a/tests/test_scene_wrapper_hygiene.py b/tests/test_scene_wrapper_hygiene.py index b2e13efb..94577095 100644 --- a/tests/test_scene_wrapper_hygiene.py +++ b/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)