From ada1f4ada15a62321fc10becb6d0fed6a7a2d148 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Fri, 22 May 2026 23:24:53 +0300 Subject: [PATCH] Fix injection performance: skip hooks when i2v off, keep tensors on GPU Two causes of 3-4x slowdown: 1. Hooks registered for blessed maps capture even when reference_injection is disabled. Now hooks only register when has_i2v is True. 2. Hidden state tensors were moved CPU on capture and GPU on inject - for 150-frame LTXAV the tensors are large (~100MB per block). Tensors now stay on the model's device throughout, eliminating the CPU<->GPU round trip overhead. Co-Authored-By: Claude Sonnet 4.6 --- ltx_enhancements.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/ltx_enhancements.py b/ltx_enhancements.py index 5b0aa58a..6e3acc22 100644 --- a/ltx_enhancements.py +++ b/ltx_enhancements.py @@ -640,10 +640,10 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re dit[("double_block", block_idx)] = replacement # --- Technique 5 (forward hooks): capture + inject via nn.Module hooks --- - # patches_replace["dit"] reaches transformer_options but LTXAV blocks never call it. - # Forward hooks on the actual nn.Modules are the only reliable mechanism. + # Only register hooks when i2v injection is active. Hooks on large LTXAV hidden + # state tensors are expensive - skip entirely when reference_injection is off. hook_handles = [] - if capture_buf is not None: + if capture_buf is not None and has_i2v: try: diff = getattr(getattr(model, "model", None), "diffusion_model", None) if diff is not None: @@ -671,12 +671,12 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re return out, out def _hook(module, inp, out): - # Capture only at mid-sigma + # Capture only at mid-sigma - keep on same device as model if 0.85 <= s_state[0] <= 0.95: try: t, _ = _extract_tensor(out) if isinstance(t, torch.Tensor): - buf[block_idx] = t.detach().cpu() + buf[block_idx] = t.detach() except Exception: pass # Inject when lazy has a tensor @@ -690,7 +690,7 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re t, _ = _extract_tensor(out) if not isinstance(t, torch.Tensor): return None - b = lazy_ref.tensor.to(device=t.device, dtype=t.dtype) + b = lazy_ref.tensor.to(device=t.device, dtype=t.dtype, non_blocking=True) if b.shape != t.shape: return None injected = t.lerp(b, effective) @@ -776,7 +776,7 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re _ref_extracted[0] = True for idx, lazy in _lazy.items(): if idx in _cap_buf: - lazy.set(_cap_buf[idx]) + lazy.set(_cap_buf[idx]) # tensor already on GPU from capture hook return result