nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix injection performance: skip hooks when i2v off, keep tensors on GPU

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-05-22 23:24:53 +03:00
co-authored by Claude Sonnet 4.6
parent 770779b0a5
commit ada1f4ada1
1 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ltx_enhancements.py
+7 -7
View File
@@ -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