From bf6cb80f9fd65e0f211178d89a52f5571f344f93 Mon Sep 17 00:00:00 2001
From: digital-garbage <133010426+digital-garbage@users.noreply.github.com>
Date: Fri, 22 May 2026 22:44:17 +0300
Subject: [PATCH] Fix diagnostic performance and report timing
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Performance: hooks now only capture when sigma in [0.85, 0.95] - one
step per scene instead of every step. Eliminates the 2x overhead and
5x slowdown on step 6.

Report timing: fires in sigma tracker when sigma < 0.5 AND scene is
snapped AND ≥2 scenes captured. Covers the last sigma=0.421875 step
which was below the old hook_remover threshold of 0.05.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 ltx_enhancements.py | 40 +++++++++++++++++++++++++++-------------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/ltx_enhancements.py b/ltx_enhancements.py
index 2ce7ea43..47ba95ba 100644
--- a/ltx_enhancements.py
+++ b/ltx_enhancements.py
@@ -672,7 +672,7 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
                             inj_gate = (0.2, 0.5) if is_anchor else (0.35, 0.80)
                             lazy = lazy_injects.get(idx)
 
-                            def _make_hook(block_idx, buf, lazy_ref, strength, gate, s_state):
+                            def _make_hook(block_idx, buf, lazy_ref, strength, gate, s_state, scene_cur=_scene_cur_buf):
                                 def _extract_tensor(out):
                                     if isinstance(out, dict):
                                         return out.get("img") or out.get("hidden_states") or next(iter(out.values()), None), out
@@ -681,13 +681,16 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
                                     return out, out
 
                                 def _hook(module, inp, out):
-                                    # Capture
-                                    try:
-                                        t, _ = _extract_tensor(out)
-                                        if isinstance(t, torch.Tensor):
-                                            buf[block_idx] = t.detach().cpu()
-                                    except Exception:
-                                        pass
+                                    # Capture only at mid-sigma (avoid per-step tensor copies)
+                                    if 0.85 <= s_state[0] <= 0.95:
+                                        try:
+                                            t, _ = _extract_tensor(out)
+                                            if isinstance(t, torch.Tensor):
+                                                buf[block_idx] = t.detach().cpu()
+                                                if scene_cur is not None:
+                                                    scene_cur[0][block_idx] = buf[block_idx]
+                                        except Exception:
+                                            pass
                                     # Inject (only when lazy has a tensor)
                                     if lazy_ref is None or lazy_ref.tensor is None:
                                         return None
@@ -717,11 +720,15 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
                                 _make_hook(idx, capture_buf, lazy, inj_strength, inj_gate, sigma_state)
                             )
                             hook_handles.append(h)
-                    # Register capture-only diagnostic hooks on extra candidate blocks
+                    # Register capture-only diagnostic hooks on extra candidate blocks.
+                    # Gate by sigma to avoid capturing (and copying tensors) on every step.
                     for idx in _diag_blocks - blocks_to_hook:
                         if idx < len(block_list):
-                            def _make_diag_hook(block_idx, scene_cur):
+                            def _make_diag_hook(block_idx, scene_cur, s_state):
                                 def _hook(module, inp, out):
+                                    # Only capture in the mid-sigma window once per scene
+                                    if not (0.85 <= s_state[0] <= 0.95):
+                                        return None
                                     try:
                                         if isinstance(out, dict):
                                             t = out.get("img") or out.get("hidden_states") or next(iter(out.values()), None)
@@ -733,8 +740,9 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
                                             scene_cur[0][block_idx] = t.detach().cpu()
                                     except Exception:
                                         pass
+                                    return None
                                 return _hook
-                            h = block_list[idx].register_forward_hook(_make_diag_hook(idx, _scene_cur_buf))
+                            h = block_list[idx].register_forward_hook(_make_diag_hook(idx, _scene_cur_buf, sigma_state))
                             hook_handles.append(h)
 
                     # Also redirect the main capture hooks to write into scene_cur_buf too
@@ -833,8 +841,14 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
 
             # Per-scene diagnostic snapshot at mid-sigma
             if scene_cur is not None and not _scene_snapped[0] and sigma < 0.95:
-                _scene_snapped[0] = True
-                scene_cur[0] = {k: v for k, v in scene_cur[0].items()}  # snapshot copy
+                _scene_snapped[0] = True  # hooks already wrote into scene_cur[0] at this sigma
+
+            # Fire similarity report at end of each scene (sigma low) when ≥2 scenes captured
+            if scene_captures is not None and sigma < 0.5 and _scene_snapped[0]:
+                if scene_cur is not None and scene_cur[0] and (not scene_captures or scene_captures[-1] is not scene_cur[0]):
+                    scene_captures.append(dict(scene_cur[0]))
+                if len(scene_captures) >= 2:
+                    _report_block_similarity(scene_captures)
 
             return result