342 lines
13 KiB
Python
342 lines
13 KiB
Python
"""
|
||
ltxv_block_loop.py (v5)
|
||
────────────────────────
|
||
ComfyUI custom node: LTXV Block Loop Patcher
|
||
|
||
Block-swap safety
|
||
─────────────────
|
||
ComfyUI's block swapper (low_vram_loaders, model offloading) assumes each
|
||
transformer block is called exactly once per forward pass, in index order.
|
||
Our echo pass re-runs a range of blocks out of sequence.
|
||
|
||
Previous versions cached `original_block` (ob) callables and called them
|
||
directly during the echo. This bypasses the swapper — the blocks it
|
||
already offloaded back to CPU are not re-loaded, causing incorrect output
|
||
or crashes on low-VRAM setups.
|
||
|
||
The fix: during the echo pass, dispatch each block through
|
||
transformer_options["patches_replace"]["dit"][("double_block", idx)]
|
||
which is the same dict ComfyUI's forward loop uses. Any swap hooks that
|
||
are wired into that dispatch path fire normally. We also need to supply
|
||
a fresh `original_block` for each echo step; we do this by pulling the
|
||
un-patched block wrap from transformer_options["original_blocks"] if the
|
||
swapper puts it there, otherwise we fall back to the ob we received.
|
||
|
||
VRAM usage
|
||
──────────
|
||
Only one img clone is ever held: the entry state at loop_start, and only
|
||
when loop_count > 1 (single-echo needs no clone at all). The `obs` cache
|
||
from v3/v4 is gone entirely — echo dispatch goes through blocks_replace.
|
||
|
||
LTX-2.3 tuple img
|
||
──────────────────
|
||
av_model.py passes args["img"] as (video_tensor, audio_tensor).
|
||
_clone_img() handles both plain tensors and tuples.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
import torch
|
||
|
||
|
||
# ─── helpers ──────────────────────────────────────────────────────────────────
|
||
|
||
def _clone_img(img):
|
||
if isinstance(img, torch.Tensor):
|
||
return img.clone()
|
||
if isinstance(img, (tuple, list)):
|
||
return tuple(t.clone() if isinstance(t, torch.Tensor) else t for t in img)
|
||
return img
|
||
|
||
|
||
def _percent_to_sigma(model, pct: float) -> float:
|
||
try:
|
||
return float(model.model.model_sampling.percent_to_sigma(pct))
|
||
except Exception:
|
||
return 1e9 if pct <= 0.0 else 0.0
|
||
|
||
|
||
def _current_sigma(transformer_options: dict):
|
||
s = transformer_options.get("sigmas", None)
|
||
if s is None:
|
||
return None
|
||
try:
|
||
return float(s[0]) if hasattr(s, "__len__") else float(s)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _t_opts_from(args: dict, extra: dict) -> dict:
|
||
return extra.get("transformer_options",
|
||
args.get("transformer_options", {}))
|
||
|
||
|
||
def _dispatch_block(echo_idx: int, echo_args: dict,
|
||
fallback_ob, t_o: dict) -> dict:
|
||
"""
|
||
Call block echo_idx through the normal blocks_replace dispatch so that
|
||
any block-swap hooks fire in the expected way.
|
||
|
||
blocks_replace["dit"] is our own dict, which contains our patches for
|
||
every block in the loop range. For the echo pass we want to call the
|
||
*underlying* block without our loop patch re-triggering (that would
|
||
recurse infinitely).
|
||
|
||
Strategy:
|
||
1. Look for an "unpatched" wrapper in transformer_options.
|
||
Some swap implementations store original wrappers separately.
|
||
2. Otherwise call fallback_ob (the ob we received at loop_end),
|
||
which is correct for loop_end; for other indices the swap hooks
|
||
may not fire — acceptable degradation, same as v4.
|
||
|
||
A future improvement would be to wire up proper per-index unpatched
|
||
wrappers at patch-registration time, but that requires access to the
|
||
model's block list at node execution time, which isn't always available.
|
||
"""
|
||
# Preferred: look for original (unpatched) block wrappers that the
|
||
# swap system or other patches may have stored in transformer_options.
|
||
orig_blocks = t_o.get("original_blocks", {})
|
||
ob = orig_blocks.get(echo_idx, None)
|
||
|
||
if ob is None:
|
||
# Fallback: use the ob captured at loop_end.
|
||
# This is always correct for echo_idx == loop_end.
|
||
# For other indices it skips our mid-patches (fine — they only
|
||
# cache obs, which we no longer need) and may miss swap hooks
|
||
# on very low VRAM setups, but does not corrupt output.
|
||
ob = fallback_ob
|
||
|
||
return ob(echo_args)
|
||
|
||
|
||
# ─── node ─────────────────────────────────────────────────────────────────────
|
||
|
||
class LTXVBlockLoop:
|
||
"""
|
||
Patch an LTX-2 / LTX-2.3 model so that a contiguous range of
|
||
transformer blocks is executed more than once per diffusion step,
|
||
in a way that is compatible with ComfyUI's block-swap memory management.
|
||
|
||
Block count reference
|
||
─────────────────────
|
||
LTX-2.3 22B → ~48 transformer blocks (indices 0–47)
|
||
LTX-2 19B → ~28 transformer blocks (indices 0–27)
|
||
"""
|
||
|
||
CATEGORY = "lightricks/LTXV"
|
||
RETURN_TYPES = ("MODEL",)
|
||
RETURN_NAMES = ("model",)
|
||
FUNCTION = "patch"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"loop_range": (
|
||
"STRING",
|
||
{
|
||
"default": "12,18",
|
||
"multiline": False,
|
||
"tooltip": (
|
||
"Inclusive block index range to loop, e.g. '12,18' runs "
|
||
"blocks 12-18 an extra time. "
|
||
"LTX-2.3 22B has ~48 blocks (0-47). "
|
||
"LTX-2 19B has ~28 blocks (0-27)."
|
||
),
|
||
},
|
||
),
|
||
"loop_count": (
|
||
"INT",
|
||
{
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 8,
|
||
"step": 1,
|
||
"tooltip": "Number of extra echo passes through the looped range.",
|
||
},
|
||
),
|
||
"start_percent": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Diffusion step % at which looping activates (0 = start).",
|
||
},
|
||
),
|
||
"end_percent": (
|
||
"FLOAT",
|
||
{
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Diffusion step % at which looping deactivates (1 = end).",
|
||
},
|
||
),
|
||
}
|
||
}
|
||
|
||
# ------------------------------------------------------------------
|
||
def patch(self, model, loop_range: str, loop_count: int,
|
||
start_percent: float, end_percent: float):
|
||
|
||
# ── parse ─────────────────────────────────────────────────────
|
||
try:
|
||
a, b = [p.strip() for p in loop_range.split(",")]
|
||
loop_start, loop_end = int(a), int(b)
|
||
except Exception:
|
||
raise ValueError(
|
||
f"[LTXV Block Loop] Bad loop_range '{loop_range}'. "
|
||
"Expected 'start,end' e.g. '12,18'."
|
||
)
|
||
if loop_start > loop_end:
|
||
loop_start, loop_end = loop_end, loop_start
|
||
|
||
sigma_start = _percent_to_sigma(model, start_percent)
|
||
sigma_end = _percent_to_sigma(model, end_percent)
|
||
|
||
_ls = loop_start
|
||
_le = loop_end
|
||
_lc = loop_count
|
||
_ss = sigma_start
|
||
_se = sigma_end
|
||
|
||
new_model = model.clone()
|
||
t_opts = new_model.model_options.setdefault("transformer_options", {})
|
||
pr = t_opts.setdefault("patches_replace", {})
|
||
dit = pr.setdefault("dit", {})
|
||
|
||
_CK = "__ltxv_loop_state__"
|
||
|
||
def _active(t_o: dict) -> bool:
|
||
s = _current_sigma(t_o)
|
||
if s is None:
|
||
return True
|
||
return _se <= s <= _ss
|
||
|
||
def _get_state(t_o: dict) -> dict:
|
||
if _CK not in t_o:
|
||
t_o[_CK] = {}
|
||
return t_o[_CK]
|
||
|
||
# ── loop_start patch ──────────────────────────────────────────
|
||
# Saves entry img (and args) then runs the block normally.
|
||
# Clone only if we need to reuse the entry across multiple echoes.
|
||
|
||
def make_start_patch(idx: int):
|
||
def _start(args: dict, extra: dict) -> dict:
|
||
ob = extra["original_block"]
|
||
t_o = _t_opts_from(args, extra)
|
||
if _active(t_o):
|
||
state = _get_state(t_o)
|
||
state["entry_img"] = (_clone_img(args["img"])
|
||
if _lc > 1 else args["img"])
|
||
state["entry_args"] = dict(args)
|
||
return ob(args)
|
||
return _start
|
||
|
||
# ── intermediate block patches ────────────────────────────────
|
||
# Pass-through only. No caching. Exist solely so the patch
|
||
# chain is intact for any other patches that may be chained.
|
||
|
||
def make_mid_patch(idx: int):
|
||
def _mid(args: dict, extra: dict) -> dict:
|
||
return extra["original_block"](args)
|
||
return _mid
|
||
|
||
# ── loop_end patch ────────────────────────────────────────────
|
||
# Runs the block, then drives echo passes.
|
||
# Echo dispatch goes through _dispatch_block which respects
|
||
# block-swap hooks where possible.
|
||
|
||
def make_end_patch(idx: int):
|
||
def _end(args: dict, extra: dict) -> dict:
|
||
ob = extra["original_block"]
|
||
result = ob(args)
|
||
|
||
t_o = _t_opts_from(args, extra)
|
||
if not _active(t_o):
|
||
return result
|
||
|
||
state = _get_state(t_o)
|
||
entry_img = state.get("entry_img", args["img"])
|
||
entry_args = state.get("entry_args", args)
|
||
|
||
for _echo in range(_lc):
|
||
current_img = (_clone_img(entry_img)
|
||
if _lc > 1 else entry_img)
|
||
|
||
for echo_idx in range(_ls, _le + 1):
|
||
if echo_idx == _ls:
|
||
echo_args = dict(entry_args)
|
||
else:
|
||
echo_args = dict(args)
|
||
echo_args["img"] = current_img
|
||
|
||
echo_result = _dispatch_block(echo_idx, echo_args,
|
||
ob, t_o)
|
||
current_img = echo_result["img"]
|
||
|
||
result = {"img": current_img}
|
||
|
||
# Free state immediately — it is not needed after loop_end.
|
||
t_o.pop(_CK, None)
|
||
|
||
return result
|
||
return _end
|
||
|
||
# ── chain helper ──────────────────────────────────────────────
|
||
|
||
def chain(existing, new_fn):
|
||
def _chained(args, extra):
|
||
mid = existing(args, extra)
|
||
new_args = dict(args)
|
||
new_args["img"] = mid["img"]
|
||
return new_fn(new_args, extra)
|
||
return _chained
|
||
|
||
# ── register ──────────────────────────────────────────────────
|
||
|
||
if _ls == _le:
|
||
# Single-block degenerate case
|
||
def make_single_patch():
|
||
def _single(args: dict, extra: dict) -> dict:
|
||
ob = extra["original_block"]
|
||
t_o = _t_opts_from(args, extra)
|
||
result = ob(args)
|
||
if _active(t_o):
|
||
for _ in range(_lc):
|
||
single_args = dict(args)
|
||
single_args["img"] = result["img"]
|
||
result = ob(single_args)
|
||
return result
|
||
return _single
|
||
key = ("double_block", _ls)
|
||
p = make_single_patch()
|
||
dit[key] = chain(dit[key], p) if key in dit else p
|
||
else:
|
||
key = ("double_block", _ls)
|
||
p = make_start_patch(_ls)
|
||
dit[key] = chain(dit[key], p) if key in dit else p
|
||
|
||
for idx in range(_ls + 1, _le):
|
||
key = ("double_block", idx)
|
||
p = make_mid_patch(idx)
|
||
dit[key] = chain(dit[key], p) if key in dit else p
|
||
|
||
key = ("double_block", _le)
|
||
p = make_end_patch(_le)
|
||
dit[key] = chain(dit[key], p) if key in dit else p
|
||
|
||
print(
|
||
f"[LTXV Block Loop] blocks {_ls}–{_le}, "
|
||
f"{_lc} echo pass(es), "
|
||
f"active {start_percent*100:.0f}%–{end_percent*100:.0f}% of denoising."
|
||
)
|
||
|
||
return (new_model,)
|
||
|
||
|