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

feat: LTX temporal style (frame-rate styles, pulse, rapid, Mobius loop) as a model module; unet-wrapper guard

Browse Source 
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-10-03 10:52:41 +03:00
co-authored by Claude Sonnet 5.5
parent a51faa729b
commit bc8631f683
5 changed files with 503 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/patching.py
+2
View File
@@ -260,6 +260,8 @@ NEUTRAL = {
"set_model_sampler_post_cfg_function": lambda args: args["denoised"],
"set_model_sampler_cfg_function": lambda args: args["input"] - (
args["uncond_denoised"] + (args["cond_denoised"] - args["uncond_denoised"]) * args["cond_scale"]),
"set_model_unet_function_wrapper": lambda apply_fn, args: apply_fn(args["input"], args["timestep"],
**args.get("c", {})),
"set_model_denoise_mask_function": lambda sigma, denoise_mask, extra_options=None: denoise_mask,
}
modules/sampling/temporal_style/__init__.py
+61
View File
@@ -0,0 +1,61 @@
"""Temporal style: how motion and time feel, on an LTX model. One mechanism, chosen by name.
Everything but `loop` tells the model a different frame rate than the clip's (see style.py);
`loop` rolls the latent in time so the clip's end joins its start (loop.py).
Not here: v4's `auto`, which read each scene's prompt and chose per scene. That choice is a
conditioning-side decision (the prompt lives there); it is still to be ported with the scene-chain
work, and a style set here applies to every scene of the run alike.
"""
from ..._core import patching
from . import loop as _loop
from . import style
ID = "temporal_style"
TITLE = "Temporal style"
MOUNT = "generation.sampling"
STAGE = "sampling"
CATEGORY = "sampling"
STATUS = "experimental"
REQUIRES = ["ltx_av"]
STYLES = ["natural", "accelerate", "decelerate", "freeze", "pulse",
"rapid_start", "rapid_end", "rapid_start_end", "loop"]
SETTINGS = {
"style": {
"type": "enum", "default": "natural", "label": "Temporal style",
"options": [{"value": v, "label": v.replace("_", " ")} for v in STYLES],
"hint": "How motion and time feel. accelerate/freeze hold the motion, decelerate/rapid make it punchier, "
"pulse repeats ease-downs, loop joins the end to the start.",
},
}
def install(patcher, values, key):
chosen = values.get("style", "natural")
old = patcher.model_options.get("model_function_wrapper")
state = {}
if chosen in style.MULT:
mult = style.MULT[chosen]
wrapper = style.make_wrapper(old, lambda args: mult)
note = f"{chosen}: frame rate x{mult:g}"
elif chosen == "pulse":
wrapper = style.make_wrapper(old, lambda args: style.pulse_mult(style.progress_of(args, state)))
note = "pulse"
elif chosen in ("rapid_start", "rapid_end", "rapid_start_end"):
wrapper = style.make_wrapper(old, lambda args: style.rapid_mult(style.progress_of(args, state), chosen))
note = chosen.replace("_", " ")
elif chosen == "loop":
wrapper = _loop.make_loop_temporal_wrapper(old)
note = "loop: the clip's end joins its start"
else:
return None
patcher.set_model_unet_function_wrapper(patching.tag(wrapper, key))
installed = patcher.model_options["model_function_wrapper"]
installed._funpack_prev = old # so stripping ours puts the earlier one back
return note
PROVIDES = {"modifier": install}
modules/sampling/temporal_style/loop.py
+231
View File
@@ -0,0 +1,231 @@
"""Loop: seamless looping by rolling the latent in time (Mobius, arXiv:2502.20307).
On eligible denoise steps the latent is cyclically rolled along its time axis before the forward pass
and the prediction rolled back after, so the model meets the clip's seam (last frame | first frame) as an
interior pair at a step-varying position and smooths it like any other cut. No training, no extra
forward passes.
"""
import torch
from ..._core import log
# Seamless looping without training or extra forward passes: on eligible denoise steps
# the latent is cyclically rolled along its temporal dim before the forward and the
# prediction unrolled after. The model then repeatedly sees the video's seam
# (last frame | first frame) as an INTERIOR frame pair at a step-varying position and
# smooths it like any other cut — after enough steps the sequence is consistent under
# rotation, i.e. it loops. WanVideoWrapper's "Loop Args" is the same trick for WAN.
#
# Roll starts below the near-noise plateau: content is ~pure noise up there, so there is
# nothing to smooth. Appended guide frames (carry_i2v_guides / mid-scene guide / JoyAI memory /
# custom stacks) live at the TAIL of the same T range, pinned to absolute positions by
# keyframe_idxs — so only the content region [0, T - tail) rolls and the tail (plus
# keyframe_idxs and attention entries, which reference tail tokens) stays canonical.
# Guides keep informing every rolled position (their influence smears across the cycle),
# which is exactly what identity/style references should do in a loop; they are never
# rendered as frames themselves.
LOOP_ROLL_MAX_SIGMA = 0.95
LOOP_ROLL_MIN_FRAMES = 4 # below this many CONTENT latent frames there is nothing to loop
def _van_der_corput(n):
"""Base-2 bit-reversal sequence in (0,1): 0.5, 0.25, 0.75, 0.125, 0.625… — every
prefix covers the unit interval near-uniformly, so however many steps end up
eligible, the seam visits well-spread temporal positions."""
v, denom = 0.0, 1.0
n = int(n)
while n:
denom *= 2.0
v += (n & 1) / denom
n >>= 1
return v
def _loop_stream_shapes(args):
"""Per-stream latent shapes for a packed [B,1,N] input, or None when the layout
can't be trusted. Mirrors the samplers' packed-latent handling (video first, then
audio; both streams keep their temporal dim at index 2)."""
c = args.get("c") or {}
shapes = c.get("latent_shapes")
if hasattr(shapes, "cond"):
shapes = shapes.cond
x = args.get("input")
if x is None:
return None
if not shapes:
# Single-stream (plain LTXV) fed unpacked: treat the input itself as one stream.
return [tuple(int(d) for d in x.shape)] if x.ndim == 5 else None
try:
shapes = [tuple(int(d) for d in s) for s in shapes]
import math as _math
if sum(_math.prod(s[1:]) for s in shapes) != int(x.shape[-1]):
return None # packed layout doesn't match — don't risk rolling blind
except (TypeError, ValueError):
return None
return shapes
def _loop_content_shift(frac, t, tail):
"""Frame shift for a stream with `t` total frames of which the last `tail` are pinned
guide frames. 0 when the content region is too small to roll."""
tc = int(t) - max(0, int(tail))
if tc < 2:
return 0, tc
return int(round(frac * tc)) % tc, tc
def _loop_roll_stream(stream, frac, direction, tail):
"""Roll only the content region [0, T-tail) of one stream's temporal dim (dim 2);
the pinned guide tail stays canonical. Returns the input unchanged when there is
nothing to roll."""
shift, tc = _loop_content_shift(frac, stream.shape[2], tail)
if not shift:
return stream
out = stream.clone()
out[:, :, :tc] = torch.roll(stream[:, :, :tc], shifts=direction * shift, dims=2)
return out
def _loop_roll_packed(x, shapes, frac, direction, tails=None):
"""Cyclically roll every stream's content region of a packed [B,1,N] tensor by
round(frac * T_content) frames (video and audio each by their own content length, so
the two stay time-aligned). `tails` gives per-stream pinned tail-frame counts aligned
with `shapes` (appended guides / audio memory — left canonical). direction=+1 rolls,
-1 unrolls. Returns a new tensor."""
import math as _math
if x.ndim == 5: # unpacked single stream
return _loop_roll_stream(x, frac, direction, (tails or [0])[0])
out = x.clone()
off = 0
for i, dims in enumerate(shapes):
sz = _math.prod(dims[1:])
if len(dims) >= 3 and dims[2] >= 2:
stream = x[..., off:off + sz].reshape([x.shape[0]] + list(dims[1:]))
rolled = _loop_roll_stream(stream, frac, direction, (tails or [])[i] if tails else 0)
if rolled is not stream:
out[..., off:off + sz] = rolled.reshape(x.shape[0], 1, sz)
off += sz
return out
def _loop_roll_mask(mask, frac, direction, tail=0):
"""Roll a denoise mask ([B,1,T,H,W] video / [B,C,T,F] audio) in step with its stream,
leaving the pinned tail region canonical like the stream itself."""
if not isinstance(mask, torch.Tensor) or mask.ndim < 3 or mask.shape[2] < 2:
return mask
return _loop_roll_stream(mask, frac, direction, tail)
def _loop_video_tail_frames(c, video_shape):
"""Number of appended guide latent frames pinned at the video stream's tail. 0 when
the call carries no guides; None when a guide context exists but can't be parsed
(caller then leaves the call canonical rather than risking a corrupting roll).
keyframe_idxs is per-TOKEN — frames = tokens / (H*W), same math as comfy's
get_keyframe_idxs."""
kf = c.get("keyframe_idxs")
if kf is None:
# Attention entries without keyframe_idxs shouldn't happen; treat as unparseable.
return None if c.get("guide_attention_entries") else 0
try:
if len(video_shape) != 5 or not hasattr(kf, "shape") or len(kf.shape) < 3:
return None
tokens_per_frame = int(video_shape[-2]) * int(video_shape[-1])
n = int(kf.shape[2]) // max(1, tokens_per_frame)
return n if 0 < n < int(video_shape[2]) else None
except Exception:
return None
def _loop_audio_tail_frames(mask):
"""Appended audio-memory frames at the tail of the audio stream = the trailing run of
all-zero frames in its denoise mask (JoyAI paired audio memory rides mask=0; genuine
audio content is never fully masked). 0 when there is no mask or no zero tail."""
if not isinstance(mask, torch.Tensor) or mask.ndim < 3 or mask.shape[2] < 2:
return 0
try:
other_dims = [d for d in range(mask.ndim) if d != 2]
per_frame = mask.float().abs().amax(dim=other_dims) # [T]
tail = 0
for v in reversed(per_frame.tolist()):
if v > 1e-4:
break
tail += 1
return tail if tail < mask.shape[2] else 0
except Exception:
return 0
def make_loop_temporal_wrapper(old_wrapper):
"""Build the loop-style model_function_wrapper. Installed INNERMOST (closest to
apply_model): prediction-modifying wrappers layered above it (dynashift, output
guidance, …) must see canonical-orientation inputs and outputs — the roll exists
only for the duration of the base forward. The per-step shift follows a van der
Corput sequence, reset whenever sigma jumps back up (a new scene/run)."""
state = {"count": 0, "last_sigma": None, "logged": False}
def _loop_wrapper(apply_fn, args, _old=old_wrapper):
def _call(a):
if _old is not None:
return _old(apply_fn, a)
return apply_fn(a["input"], a["timestep"], **a.get("c", {}))
try:
ts = args.get("timestep")
sigma = float(ts.max().item()) if ts is not None else 1.0
except Exception:
sigma = 1.0
if state["last_sigma"] is not None and sigma > state["last_sigma"] + 1e-4:
state["count"] = 0 # sigma went back up: new scene/run
state["last_sigma"] = sigma
c = args.get("c") or {}
if sigma > LOOP_ROLL_MAX_SIGMA:
return _call(args)
shapes = _loop_stream_shapes(args)
if not shapes:
return _call(args)
# Per-stream pinned tails: appended guide frames (video, counted from the
# per-token keyframe_idxs) and appended audio memory (trailing mask=0 run).
# Only the content region in front of each tail rolls; the tail — and the
# keyframe_idxs / attention entries that reference it — stays canonical.
video_shape = next((s for s in shapes if len(s) == 5), shapes[0])
v_tail = _loop_video_tail_frames(c, video_shape)
if v_tail is None:
return _call(args) # guide context we can't parse: never risk a bad roll
tails = [
(v_tail if len(s) == 5 else _loop_audio_tail_frames(c.get("audio_denoise_mask")))
for s in shapes
]
t_content = int(video_shape[2]) - v_tail if len(video_shape) >= 3 else 0
if t_content < LOOP_ROLL_MIN_FRAMES:
return _call(args)
state["count"] += 1
frac = _van_der_corput(state["count"])
if not state["logged"]:
state["logged"] = True
tail_note = f", {v_tail} guide frame(s) pinned" if v_tail else ""
log.once("temporal_style:loop", log.INFO, "FunPack Temporal style", f"loop: Mobius latent roll active (T={t_content}{tail_note})")
try:
new_c = dict(c)
if isinstance(new_c.get("denoise_mask"), torch.Tensor):
new_c["denoise_mask"] = _loop_roll_mask(new_c["denoise_mask"], frac, 1, tail=v_tail)
if isinstance(new_c.get("audio_denoise_mask"), torch.Tensor):
new_c["audio_denoise_mask"] = _loop_roll_mask(
new_c["audio_denoise_mask"], frac, 1,
tail=_loop_audio_tail_frames(new_c["audio_denoise_mask"]))
rolled = dict(args)
rolled["input"] = _loop_roll_packed(args["input"], shapes, frac, 1, tails=tails)
rolled["c"] = new_c
out = _call(rolled)
return _loop_roll_packed(out, shapes, frac, -1, tails=tails)
except Exception as e:
if state.get("roll_error") is None:
state["roll_error"] = True
log.once("temporal_style:loop_failed", log.ALERT, "FunPack Temporal style", f"the loop roll failed ({type(e).__name__}: {e}); this call ran without it")
return _call(args)
return _loop_wrapper
modules/sampling/temporal_style/style.py
+81
View File
@@ -0,0 +1,81 @@
"""How time is felt: the frame rate the model THINKS it is making.
LTX positions every frame in time from the frame rate in its conditioning. Telling it a higher rate
for the same frames means smaller steps between frames (smoother, more held motion); a lower one means
larger steps (punchier). The names track the clock, not the subject.
"""
# frame_rate multiplier per style; >1 holds, <1 punches.
MULT = {"accelerate": 1.35, "decelerate": 0.72, "freeze": 2.0}
# Pulse: repeated ease-down segments, each starting punchy and easing to a hold.
PULSE_SEGMENTS, PULSE_PEAK, PULSE_FLOOR = 3, 0.88, 1.65
# Rapid start/end: a punchy multiplier over the first and/or last slice of the denoise, natural elsewhere.
RAPID_MULT, RAPID_FRACTION = 0.65, 0.35
def _ease(t):
t = max(0.0, min(1.0, float(t)))
return t * t * (3.0 - 2.0 * t)
def pulse_mult(progress, segments=PULSE_SEGMENTS, peak=PULSE_PEAK, floor=PULSE_FLOOR):
progress = max(0.0, min(1.0, float(progress)))
index = min(int(progress * segments), segments - 1)
return peak + (floor - peak) * _ease(progress * segments - index)
def rapid_mult(progress, mode, fraction=RAPID_FRACTION, mult=RAPID_MULT):
fraction = max(1e-3, min(0.5, float(fraction)))
progress = max(0.0, min(1.0, float(progress)))
if mode in ("rapid_start", "rapid_start_end") and progress < fraction:
return mult + (1.0 - mult) * _ease(progress / fraction)
if mode in ("rapid_end", "rapid_start_end") and progress > 1.0 - fraction:
return mult + (1.0 - mult) * _ease((1.0 - progress) / fraction)
return 1.0
def scale_frame_rate(args, mult):
"""`args` with the conditioning's frame_rate times `mult`; unchanged when it has none."""
c = args.get("c")
if not (isinstance(c, dict) and "frame_rate" in c):
return args
cond = c["frame_rate"]
if not hasattr(cond, "cond"):
return args
scaled = dict(c)
scaled["frame_rate"] = type(cond)(float(cond.cond) * float(mult))
return {**args, "c": scaled}
def progress_of(args, state):
"""How far through the denoise this call is, 0..1: from the schedule when the sampler published
it, else from how far sigma has fallen since this run's first call."""
from ..._core import dit_hooks
where = dit_hooks.current_step((args.get("c") or {}).get("transformer_options"))
if where is not None:
index, total = where
return index / max(1, total - 1) if total > 1 else 1.0
ts = args.get("timestep")
try:
sigma = float(ts.max().item()) if ts is not None else 1.0
except Exception: # noqa: BLE001
sigma = 1.0
if state.get("start") is None or sigma > state.get("last", 0.0) + 1e-4:
state["start"] = max(sigma, 1e-6) # sigma went back up: a new run
state["last"] = sigma
return max(0.0, min(1.0, 1.0 - sigma / state["start"]))
def make_wrapper(old, mult_for):
"""A model_function_wrapper that scales the frame rate by `mult_for(args)`, then runs `old` (the
wrapper that was there) or the model."""
def wrapper(apply_fn, args):
mult = mult_for(args)
if abs(mult - 1.0) > 1e-3:
args = scale_frame_rate(args, mult)
if old is not None:
return old(apply_fn, args)
return apply_fn(args["input"], args["timestep"], **args.get("c", {}))
return wrapper
modules/sampling/temporal_style/tests/test_temporal_style.py
+128
View File
@@ -0,0 +1,128 @@
import types
import pytest
import torch
from core import patching
KEY = "funpack.temporal_style"
class Cond:
def __init__(self, cond):
self.cond = cond
def _install(tiny, style, old=None):
from modules.sampling.temporal_style import install
p = tiny.patcher.clone()
if old is not None:
p.set_model_unet_function_wrapper(old)
g = patching.GuardedPatcher(p, KEY, patching.Dropped())
return p, install(g, {"style": style}, key=KEY)
def _call(patched, **extra):
seen = {}
def apply_fn(x, t, **c):
seen["frame_rate"] = c["frame_rate"].cond
seen["x"] = x
return x
args = {"input": torch.zeros(1, 1, 4), "timestep": torch.tensor([0.5]),
"c": {"frame_rate": Cond(25.0), **extra}}
patched.model_options["model_function_wrapper"](apply_fn, args)
return seen
def test_natural_installs_nothing(tiny_ltx):
p, note = _install(tiny_ltx, "natural")
assert note is None and "model_function_wrapper" not in p.model_options
@pytest.mark.parametrize("style, mult", [("accelerate", 1.35), ("decelerate", 0.72), ("freeze", 2.0)])
def test_the_frame_rate_the_model_sees_is_scaled(tiny_ltx, style, mult):
p, note = _install(tiny_ltx, style)
assert _call(p)["frame_rate"] == pytest.approx(25.0 * mult) and style in note
def test_pulse_and_rapid_follow_the_schedules_progress(tiny_ltx):
sched = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2, 0.0])
def at(style, step):
p, _ = _install(tiny_ltx, style)
return _call(p, transformer_options={"sample_sigmas": sched, "sigmas": sched[step:step + 1]})["frame_rate"] / 25.0
assert at("rapid_start", 0) == pytest.approx(0.65) and at("rapid_start", 2) == pytest.approx(1.0)
assert at("rapid_end", 5) == pytest.approx(0.65) and at("rapid_end", 1) == pytest.approx(1.0)
assert at("pulse", 0) == pytest.approx(0.88)
def test_an_earlier_wrapper_still_runs_and_comes_back_when_ours_is_stripped(tiny_ltx):
ran = []
def earlier(apply_fn, args):
ran.append(args["c"]["frame_rate"].cond)
return apply_fn(args["input"], args["timestep"], **args["c"])
p, _ = _install(tiny_ltx, "freeze", old=earlier)
_call(p)
assert ran == [50.0] # ours scaled, then the earlier one ran
patching.strip(p, KEY)
assert p.model_options["model_function_wrapper"] is earlier
def test_a_failing_wrapper_is_dropped_and_the_model_still_runs(tiny_ltx):
from modules.sampling.temporal_style import install
p = tiny_ltx.patcher.clone()
dropped = patching.Dropped()
install(patching.GuardedPatcher(p, KEY, dropped), {"style": "freeze"}, key=KEY)
class Broken:
@property
def cond(self):
raise ValueError("no frame rate")
ran = []
args = {"input": torch.zeros(1, 1, 4), "timestep": torch.tensor([0.5]), "c": {"frame_rate": Broken()}}
out = p.model_options["model_function_wrapper"](lambda x, t, **c: ran.append(1) or x, args)
assert ran == [1] and out is args["input"] and "funpack.temporal_style" in dropped
def _loop_args(video, audio, sigma=0.5):
packed = torch.cat([video.reshape(1, 1, -1), audio.reshape(1, 1, -1)], dim=-1)
return {"input": packed, "timestep": torch.tensor([sigma]),
"c": {"latent_shapes": [video.shape, audio.shape]}}
def test_the_loop_roll_is_undone_on_the_way_out_and_moves_the_seam_in_between(tiny_ltx):
p, note = _install(tiny_ltx, "loop")
video, audio = torch.randn(1, 8, 8, 2, 2), torch.randn(1, 4, 8, 3)
seen = {}
def apply_fn(x, t, **c):
seen["x"] = x.clone()
return x * 2.0
args = _loop_args(video, audio)
out = p.model_options["model_function_wrapper"](apply_fn, args)
assert torch.allclose(out, args["input"] * 2.0) # rolled in, rolled back out: canonical orientation
assert not torch.equal(seen["x"], args["input"]) # the model saw it rolled
assert "loop" in note
def test_the_loop_leaves_the_noisy_early_steps_and_short_clips_alone(tiny_ltx):
p, _ = _install(tiny_ltx, "loop")
seen = {}
def apply_fn(x, t, **c):
seen["x"] = x.clone()
return x
early = _loop_args(torch.randn(1, 8, 8, 2, 2), torch.randn(1, 4, 8, 3), sigma=0.99)
p.model_options["model_function_wrapper"](apply_fn, early)
assert torch.equal(seen["x"], early["input"])
short = _loop_args(torch.randn(1, 8, 2, 2, 2), torch.randn(1, 4, 2, 3))
p.model_options["model_function_wrapper"](apply_fn, short)
assert torch.equal(seen["x"], short["input"])