From bc8631f683888a4189442f322fb74c80ce449eca Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 3 Oct 2026 10:52:41 +0300 Subject: [PATCH] feat: LTX temporal style (frame-rate styles, pulse, rapid, Mobius loop) as a model module; unet-wrapper guard Co-Authored-By: Claude Sonnet 5.5 --- core/patching.py | 2 + modules/sampling/temporal_style/__init__.py | 61 +++++ modules/sampling/temporal_style/loop.py | 231 ++++++++++++++++++ modules/sampling/temporal_style/style.py | 81 ++++++ .../tests/test_temporal_style.py | 128 ++++++++++ 5 files changed, 503 insertions(+) create mode 100644 modules/sampling/temporal_style/__init__.py create mode 100644 modules/sampling/temporal_style/loop.py create mode 100644 modules/sampling/temporal_style/style.py create mode 100644 modules/sampling/temporal_style/tests/test_temporal_style.py diff --git a/core/patching.py b/core/patching.py index 15d00b6c..ac40f371 100644 --- a/core/patching.py +++ b/core/patching.py @@ -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, } diff --git a/modules/sampling/temporal_style/__init__.py b/modules/sampling/temporal_style/__init__.py new file mode 100644 index 00000000..b91734da --- /dev/null +++ b/modules/sampling/temporal_style/__init__.py @@ -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} diff --git a/modules/sampling/temporal_style/loop.py b/modules/sampling/temporal_style/loop.py new file mode 100644 index 00000000..c720c81f --- /dev/null +++ b/modules/sampling/temporal_style/loop.py @@ -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 + + diff --git a/modules/sampling/temporal_style/style.py b/modules/sampling/temporal_style/style.py new file mode 100644 index 00000000..58ddc760 --- /dev/null +++ b/modules/sampling/temporal_style/style.py @@ -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 diff --git a/modules/sampling/temporal_style/tests/test_temporal_style.py b/modules/sampling/temporal_style/tests/test_temporal_style.py new file mode 100644 index 00000000..37dcc4f0 --- /dev/null +++ b/modules/sampling/temporal_style/tests/test_temporal_style.py @@ -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"])