nodes/malkuthro-ComfyUI-Koolook
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(ltx): Soften Schedule — predictable step-fraction window + info output

Browse Source 
Replace the sigma-progress `crossover` knob (which covered a wildly different
number of steps per sigma curve — flat front-loaded curves stretched it across
~half the run, steep curves collapsed it to ~1 step) with `protect_until`: a
FRACTION OF THE STEPS, resolved against the sampler's ACTUAL schedule at run
time via resolve_soften_window(). 0.4 now means "soften the first 40% of the
steps" on any curve / step count — no per-schedule re-tuning.

Add an optional SIGMAS input + an `info` STRING output so the resolved window
(softened steps + sigma range) is visible in-graph via a text node. A coverage
line is also logged at run time. max_soften (amplitude) is unchanged.

Pure resolve_soften_window unit-tested (6 cases incl. curve-independence);
model wiring render-validated.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
AI Machine
2026-06-25 22:15:14 +02:00
co-authored by Claude Opus 4.8
parent 2ea2a520fa
commit d3dffb0d39
2 changed files with 189 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
k_ltx_keyframe_soften_schedule.py
+147 -60
View File
@@ -14,14 +14,17 @@ to ``1 - strength`` (a value < 1). The sampler applies it EVERY step
Hard pins held from step 0 force the model to hit exact poses immediately, so a
transition between two different poses resolves as a CUT. This node schedules the
pin by sigma instead: it **softens** the keyframe pins early (high sigma), when
the model is deciding coarse motion — letting it form a smooth trajectory between
the poses instead of snapping — then **restores** them by a configurable
``crossover`` measured in denoise *progress* (``1 - sigma/sigma_max``) so the poses
sharpen as denoise finishes. Because LTX's sigma schedule is front-loaded (sigma
held high, dropping late), a SMALL progress fraction already covers the early
motion-forming steps — hence the ``crossover`` widget maps its 0-1 range onto a
small internal ``0-0.05`` progress band (see ``_CROSSOVER_SCALE``).
pin instead: it **softens** the keyframe pins early, when the model is deciding
coarse motion — letting it form a smooth trajectory between the poses — then
**restores** them so the poses sharpen as denoise finishes.
The softening window is set by ``protect_until`` — a FRACTION OF THE STEPS (0-1),
resolved against the sampler's ACTUAL sigma schedule at run time. ``0.4`` always
means "soften the first 40% of the steps", on any sigma curve / step count, so it
needs no re-tuning when the schedule changes. (This replaces the old ``crossover``
knob, which was measured in sigma-progress and therefore covered a wildly
different number of steps on different curves — flat front-loaded curves stretched
it across many steps, steep curves collapsed it to ~1.)
It adds NOTHING (no extra guide frames / poses — that approach backfires; every
guide is a pose target). It only loosens/tightens the keyframes you already have.
@@ -30,8 +33,13 @@ Audio-safe: the A/V latent mask is a ``NestedTensor((video_mask, audio_mask))``
(``LTXVConcatAVLatent``); we soften ONLY the video mask and pass the audio mask
through untouched, so lip-sync timing is never disturbed.
NOTE: the pure ``soften_gain`` ramp is unit-tested; the model wiring (the
denoise-mask function) is validated by rendering against a live LTX 2.3 model.
The node emits an ``info`` STRING describing the resolved window — wire a SIGMAS
input (the same one your sampler uses) to read the exact softened steps + sigma
range in-graph via a text node.
NOTE: the pure ``soften_gain`` ramp and ``resolve_soften_window`` are unit-tested;
the model wiring (the denoise-mask function) is validated by rendering against a
live LTX 2.3 model.
"""
from __future__ import annotations
@@ -39,11 +47,9 @@ import logging
log = logging.getLogger(__name__)
# The node widgets are normalized 0-1 sliders that map onto the *useful* internal
# range (the effect is very sensitive, so the raw values bunch up near 0). Widget
# value * SCALE = the internal value passed to soften_gain().
# max_soften widget is a normalized 0-1 slider mapping onto the *useful* internal
# amplitude range (the effect is sensitive, so raw values bunch up near 0).
_MAX_SOFTEN_SCALE = 0.2 # widget 0..1 -> internal max_soften 0..0.2
_CROSSOVER_SCALE = 0.05 # widget 0..1 -> internal crossover 0..0.05
def soften_gain(sigma, sigma_max, crossover=0.55, max_soften=0.6):
@@ -53,8 +59,10 @@ def soften_gain(sigma, sigma_max, crossover=0.55, max_soften=0.6):
(``sigma -> 0``). Returns ``max_soften`` at progress 0, smoothsteps down to 0
by ``crossover`` progress, and stays 0 afterwards (pins fully restored).
A returned gain ``g`` raises a pinned mask value ``m`` toward 1 (unpinned)
via ``m + (1 - m) * g``: g=0 keeps the exact pin, g=1 fully frees it.
``crossover`` here is an internal sigma-progress threshold; the node derives it
from the step-fraction ``protect_until`` against the live schedule
(see ``resolve_soften_window``). A returned gain ``g`` raises a pinned mask
value ``m`` toward 1 (unpinned) via ``m + (1 - m) * g``.
"""
sigma_max = float(sigma_max)
if sigma_max <= 0.0:
@@ -98,6 +106,50 @@ def soften_denoise_mask(denoise_mask, gain, nested_ctor=None):
return soften_mask_value(denoise_mask, gain)
def resolve_soften_window(sigmas, protect_until):
"""Resolve the step-fraction ``protect_until`` against an actual sigma schedule.
``sigmas`` is the sampler's sigma list (length ``n_steps + 1``; the last entry
is the final ~0 sigma). ``protect_until`` in [0, 1] is the fraction of the
STEPS to soften. Returns ``(cutoff_step, threshold_sigma, n_steps)``: the
soften ramps from full at step 0 down to 0 at ``cutoff_step``, and
``threshold_sigma`` is the sigma at that step. Reading from the actual schedule
makes "first X% of steps" mean the same thing on any sigma curve.
"""
vals = [float(s) for s in sigmas]
n = len(vals)
if n <= 1:
return 0, (vals[0] if vals else 0.0), max(0, n - 1)
n_steps = n - 1
cutoff = int(round(float(protect_until) * n_steps))
cutoff = max(0, min(cutoff, n_steps))
return cutoff, vals[cutoff], n_steps
def _format_info(max_soften_w, eff_max_soften, protect_until, window):
"""Build the human-readable ``info`` string. ``window`` is the
(cutoff, threshold, n_steps) tuple from resolve_soften_window, or None."""
head = (
"LTX Keyframe Soften\n"
f" max_soften {max_soften_w:.2f} -> {eff_max_soften:.3f} "
f"(~{eff_max_soften * 100:.0f}% unpinned at peak)"
)
if eff_max_soften <= 0.0 or protect_until <= 0.0:
return head + "\n [inactive — model passed through]"
if window is not None:
cutoff, threshold, n_steps, sigma0 = window
return (
head
+ f"\n protect_until {protect_until:.2f} -> soften steps 0-{cutoff} of {n_steps}"
+ f"\n sigma {sigma0:.3f} -> {threshold:.3f} (gain ramps max->0)"
)
return (
head
+ f"\n protect_until {protect_until:.2f} = first {protect_until * 100:.0f}% of steps"
+ "\n (wire SIGMAS to preview exact steps / sigma range)"
)
class LTXKeyframeSoftenSchedule:
"""Sigma-schedule the LTX keyframe pins: soft early (smooth motion), sharp late."""
@@ -110,45 +162,50 @@ class LTXKeyframeSoftenSchedule:
"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "EASE amount — HOW HARD the keyframe pins loosen at the "
"softening peak. Amplitude only: independent of the "
"sampler/sigma schedule, so it always means the same "
"thing. Normalized 0-1 -> internal 0-0.20. Default 0.5 "
"= ~10% unpinned at peak (a hard pin 0.20 -> 0.28). "
"0 = stock (no ease); higher = smoother but more "
"sampler/sigma schedule. Normalized 0-1 -> internal "
"0-0.20. Default 0.5 = ~10% unpinned at peak (a hard pin "
"0.20 -> 0.28). 0 = stock; higher = smoother but more "
"invented in-between motion.",
}),
"crossover": ("FLOAT", {
"default": 0.65, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "INVENTION window — HOW LONG the pins stay loose before "
"re-locking, measured in denoise sigma-progress "
"(1 - sigma/sigma_max), NOT step count. NOTE: the number "
"of steps it actually covers SHIFTS if you change the "
"sampler/sigma schedule. Normalized 0-1 -> internal "
"0-0.05; default 0.65 protects roughly the first few "
"(high-sigma) steps. LOWER = lock sooner = FEWER invented "
"in-between frames; higher = more invented motion (and "
"risk of a noisy/under-resolved image).",
"protect_until": ("FLOAT", {
"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "HOW LONG to soften, as a FRACTION OF THE STEPS (resolved "
"against the live sigma schedule). 0.4 = soften the first "
"40% of the steps, on ANY curve / step count — no "
"re-tuning when the schedule changes. LOWER = lock sooner "
"= FEWER invented in-between frames; higher = more invented "
"motion (+ risk of a noisy/under-resolved image). 0 = off.",
}),
},
"optional": {
"sigmas": ("SIGMAS", {
"tooltip": "Optional. Wire the SAME sigmas you feed the sampler to "
"preview the exact softened steps + sigma range on the "
"'info' output (use a text / show-anything node). Display "
"only — behavior always uses the sampler's actual schedule "
"at run time.",
}),
},
}
RETURN_TYPES = ("MODEL",)
RETURN_NAMES = ("model",)
RETURN_TYPES = ("MODEL", "STRING")
RETURN_NAMES = ("model", "info")
FUNCTION = "patch"
CATEGORY = "Koolook/LTX"
DESCRIPTION = (
"Smooth motion early, sharp poses late: sigma-schedule the LTX keyframe "
"pins — loosen them at high sigma so the model forms smooth motion between "
"poses, restore them by a crossover (~step 9) so the poses sharpen. Adds no "
"frames; audio mask untouched (lip-sync safe)."
"pins — loosen them early so the model forms smooth motion between poses, "
"restore them by `protect_until` (a fraction of the steps) so the poses "
"sharpen. Adds no frames; audio mask untouched (lip-sync safe). The `info` "
"output reports the resolved window (wire SIGMAS to preview it in-graph)."
)
def patch(self, model, max_soften, crossover):
def patch(self, model, max_soften, protect_until, sigmas=None):
m = model.clone()
# Widgets are normalized 0-1; map onto the sensitive internal range.
max_soften_w, crossover_w = float(max_soften), float(crossover)
max_soften = max_soften_w * _MAX_SOFTEN_SCALE
crossover = crossover_w * _CROSSOVER_SCALE
max_soften_w = float(max_soften)
eff_max_soften = max_soften_w * _MAX_SOFTEN_SCALE
protect_until = float(protect_until)
try:
sigma_max = float(m.get_model_object("model_sampling").sigma_max)
@@ -156,13 +213,23 @@ class LTXKeyframeSoftenSchedule:
log.warning("[LTXKeyframeSoftenSchedule] no model_sampling.sigma_max: %s", e)
sigma_max = 1.0
if max_soften <= 0.0:
log.info("[LTXKeyframeSoftenSchedule] max_soften=0 — model passed through.")
return (m,)
# --- info string (display only; never drives behavior) ---
preview = None
if sigmas is not None:
try:
cutoff, threshold, n_steps = resolve_soften_window(sigmas, protect_until)
preview = (cutoff, threshold, n_steps, float(sigmas[0]))
except Exception as e:
log.warning("[LTXKeyframeSoftenSchedule] could not read preview sigmas: %s", e)
info = _format_info(max_soften_w, eff_max_soften, protect_until, preview)
# Chain any pre-existing denoise_mask_function instead of clobbering it
# (mirrors how the sibling LTX patchers chain model_function_wrapper).
if eff_max_soften <= 0.0 or protect_until <= 0.0:
log.info("[LTXKeyframeSoftenSchedule] max_soften/protect_until = 0 — model passed through.")
return (m, info)
# Chain any pre-existing denoise_mask_function instead of clobbering it.
prev_fn = m.model_options.get("denoise_mask_function")
state = {"crossover": None} # internal sigma-progress, resolved once at run time
diag = [False]
warned = [False]
@@ -171,37 +238,56 @@ class LTXKeyframeSoftenSchedule:
denoise_mask = prev_fn(sigma, denoise_mask, extra_options=extra_options)
if denoise_mask is None:
return denoise_mask
# Resolve protect_until -> sigma-progress crossover from the ACTUAL
# running schedule, once. This is what makes the step-fraction mean the
# same thing on any curve.
if state["crossover"] is None:
rt = (extra_options or {}).get("sigmas")
try:
cutoff, threshold, n_steps = resolve_soften_window(rt, protect_until)
state["crossover"] = max(0.0, 1.0 - threshold / sigma_max) if sigma_max > 0 else 0.0
log.info(
"[LTXKeyframeSoftenSchedule] soften steps 0-%d of %d "
"(sigma %.3f->%.3f), max_soften=%.3f",
cutoff, n_steps, float(rt[0]), threshold, eff_max_soften,
)
except Exception as _e:
state["crossover"] = 0.0
if not warned[0]:
warned[0] = True
log.warning(
"[LTXKeyframeSoftenSchedule] could not resolve the run's "
"sigma schedule (%s); softening disabled this run.", _e,
)
cp = state["crossover"]
if cp <= 0.0:
return denoise_mask
try:
s = float(sigma.max()) if hasattr(sigma, "max") else float(sigma)
except Exception as _e:
except Exception:
s = sigma_max
if not warned[0]:
warned[0] = True
log.warning(
"[LTXKeyframeSoftenSchedule] could not read sigma (%s); "
"softening at max this step.", _e,
)
gain = soften_gain(s, sigma_max, crossover, max_soften)
gain = soften_gain(s, sigma_max, cp, eff_max_soften)
if gain <= 0.0:
return denoise_mask
if not diag[0]:
diag[0] = True
nested = getattr(denoise_mask, "is_nested", False)
log.info(
"[LTXKeyframeSoftenSchedule] active (%s): softening video keyframe "
"pins%s. max_soften=%.3f crossover=%.3f",
"[LTXKeyframeSoftenSchedule] active (%s)%s.",
"A/V" if nested else "video-only",
", audio untouched" if nested else "", max_soften, crossover,
" — audio mask untouched" if nested else "",
)
return soften_denoise_mask(denoise_mask, gain)
m.set_model_denoise_mask_function(denoise_mask_function)
log.info(
"[LTXKeyframeSoftenSchedule] patched: max_soften=%.2f (->%.3f) "
"crossover=%.2f (->%.3f) sigma_max=%.3f",
max_soften_w, max_soften, crossover_w, crossover, sigma_max,
"protect_until=%.2f sigma_max=%.3f",
max_soften_w, eff_max_soften, protect_until, sigma_max,
)
return (m,)
return (m, info)
NODE_CLASS_MAPPINGS = {"LTXKeyframeSoftenSchedule": LTXKeyframeSoftenSchedule}
@@ -210,4 +296,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {"LTXKeyframeSoftenSchedule": "LTX Keyframe Soften
__all__ = [
"NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS",
"soften_gain", "soften_mask_value", "soften_denoise_mask",
"resolve_soften_window",
]
tests/nodes/test_ltx_keyframe_soften_schedule.py
+42
View File
@@ -21,6 +21,7 @@ _spec.loader.exec_module(_m)
gain = _m.soften_gain
soften = _m.soften_mask_value
soften_mask = _m.soften_denoise_mask
resolve = _m.resolve_soften_window
SMAX = 10.0
@@ -119,3 +120,44 @@ def test_soften_mask_noop_when_gain_zero():
def test_soften_mask_noop_when_none():
assert soften_mask(None, 0.5) is None
# --- resolve_soften_window: step-fraction → schedule, curve-independent --------
def test_resolve_half_of_steps():
# 4 steps (5 sigmas incl final 0); protect 0.5 -> cutoff 2 -> threshold sigmas[2]
cutoff, thr, n = resolve([1.0, 0.9, 0.6, 0.3, 0.0], 0.5)
assert (cutoff, n) == (2, 4)
assert thr == pytest.approx(0.6)
def test_resolve_zero_protects_nothing():
# cutoff 0 -> threshold == sigma_max -> downstream crossover-progress 0 -> no soften
cutoff, thr, n = resolve([1.0, 0.9, 0.0], 0.0)
assert cutoff == 0 and thr == pytest.approx(1.0) and n == 2
def test_resolve_full_protects_all():
cutoff, thr, n = resolve([1.0, 0.9, 0.0], 1.0)
assert cutoff == n == 2 and thr == pytest.approx(0.0)
def test_resolve_same_fraction_scales_with_step_count():
# The whole point: 0.5 means "half the steps" on ANY step count.
c4, _, n4 = resolve([1.0, 0.9, 0.6, 0.3, 0.0], 0.5) # 4 steps
c16, _, n16 = resolve([1.0 - i / 16 for i in range(17)], 0.5) # 16 steps
assert (c4, n4) == (2, 4)
assert (c16, n16) == (8, 16)
def test_resolve_front_loaded_curve_reads_actual_sigma():
# flat-top 16-step curve: protect 0.4 -> step 6, which is still on the plateau
sig = [1.0] * 8 + [0.9, 0.8, 0.65, 0.5, 0.35, 0.2, 0.1, 0.05, 0.0] # 17 vals, 16 steps
cutoff, thr, n = resolve(sig, 0.4)
assert n == 16 and cutoff == 6
assert thr == pytest.approx(1.0) # step 6 is still sigma 1.0 on this curve
def test_resolve_short_or_empty():
assert resolve([1.0], 0.5) == (0, 1.0, 0)
assert resolve([], 0.5) == (0, 0.0, 0)