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

anchor_shift: the sampler must see the real schedule, not its own half of it

Browse Source 
Raised by the user: "it should read the actual sampler schedule." They were right,
and this is a bigger deal than anchor_shift.

Both FunPack samplers derive their phase boundaries from len(sigmas) — the tensor
they are handed. anchor_shift runs a scene as two sample_custom calls over slices of
one schedule, so each half re-derived every phase INSIDE the slice:

  _sample_const_rf_full (Hybrid Euler 2S, and the LTXAV path of it)
    - quality phase: late_start = pct of the SLICE. On the user's 9-step schedule at
      high_quality_pct 0.35 the quality phase starts at sigma 0.603; split at 0.482 it
      starts at 0.812 in pass 1 and 0.482 in pass 2.
    - motion-pulse anchors keyed by step index within the slice
    - eta decay + velocity-bias ratio measured against the slice's first sigma
    - the euler->Heun correction ramp restarted per call
  sample_funpack_distilled_flow
    - final-correction window: the last N steps OF THE SLICE, so pass 1 ran a "final"
      Heun correction in the middle of the trajectory
    - AB2 ramp w = i/(total-1) restarted from 0 in pass 2
    - velocity-bias sigma ratio against the slice's first sigma

So turning anchor_shift on silently re-timed the sampler's own machinery in both
halves. Present since 015f274 and a far better explanation for a visible joint or a
thin ending than anything in the shift itself.

Fix: _schedule_view(sigmas, schedule_context) -> (real schedule, step offset). The
chain sampler passes {"sigmas": <the scene's schedule>, "offset": <slice start>} via
extra_options (the same chunk-scoped mechanism the alg_guide_* options already use,
always written even as None so it cannot leak into a later chunk). Both samplers take
the kwarg; every phase-relative quantity now reads the schedule and every step index
is step_abs = i + offset. No context = a normal single call = sigmas IS the schedule,
offset 0, byte-identical behaviour.

_anchor_shift_split_sigmas also returns pass 2's index into the schedule, and its
docstring now states the boundary rule the way the user framed it: walk the schedule,
the last pre-shift step is the one whose NEXT sigma has crossed the threshold. On
1.000 0.955 0.893 0.812 0.715 0.603 0.482 0.241 0.121 0.0 with shift 0.482 and restart
0.812 that gives pass 1 = 1.000..0.482 and pass 2 = 0.812..0.0, which is now asserted
directly as a test on that schedule.

Not changed, deliberately: sample_funpack_hybrid_euler_2s's own non-CONST loop (LTXAV
always takes the RF port, so no anchor_shift run reaches it) and segmented detailing's
stage-2 tail (a separate refine pass over a finished latent with its own schedule by
design, not a slice of one).

Suite at baseline: 22 pre-existing failures, 283 passing (+4).

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-07-31 12:59:39 +03:00
co-authored by Claude Opus 5
parent 5aac3cc5c1
commit 4c2a6b6b54
2 changed files with 140 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samplers.py
+89 -26
View File
@@ -340,6 +340,24 @@ def _velocity_bias_key(refinement_key, target, x):
return (key, f"{float(target):.2f}", shape)
def _schedule_view(sigmas, schedule_context):
"""The real schedule behind a possibly-partial `sigmas`, plus this run's step offset.
anchor_shift samples a scene as two sample_custom calls over slices of one schedule.
A sampler that measures its phases from `len(sigmas)` therefore re-derives them inside
each slice — the final-correction window reopens at the end of pass 1, the AB2 ramp
restarts, the velocity-bias ratio is taken against the slice's own first sigma. The
caller passes the whole schedule and where this slice starts in it; absent (any normal
single-call run) `sigmas` IS the whole schedule and the offset is 0.
"""
if isinstance(schedule_context, dict):
full = schedule_context.get("sigmas")
if isinstance(full, torch.Tensor) and int(full.numel()) >= 2:
offset = max(0, int(schedule_context.get("offset", 0) or 0))
return full, offset
return sigmas, 0
def _sigma_ratio(sigmas, sigma):
try:
start = float(sigmas[0].item())
@@ -1109,7 +1127,8 @@ def _sample_const_rf_full(model, x, sigmas, extra_args, callback, disable,
velocity_refinement_key,
rescue_mode, rescue_threshold, rescue_strength, rescue_prompt_sig,
quality_sharpness=0.0, velocity_bias_source="mean",
normalize_strength=0.0, normalize_start_sigma=0.9):
normalize_strength=0.0, normalize_start_sigma=0.9,
schedule_context=None):
"""Full-feature rectified-flow sampler for CONST models (LTXAV).
Rectified-flow-correct port of the hybrid sampler so its features actually run on
@@ -1140,9 +1159,14 @@ def _sample_const_rf_full(model, x, sigmas, extra_args, callback, disable,
correction_blend = max(0.0, min(1.0, float(correction_blend)))
eta_final = max(0.0, min(float(eta), float(eta_final)))
# See _schedule_view: `sigmas` may be one half of an anchor_shift split, and the quality
# phase / motion-pulse anchors below are positions in the SCHEDULE, not in this call.
sched, step_offset = _schedule_view(sigmas, schedule_context)
sched_steps = max(1, int(len(sched)) - 1)
if not motion_pulse_steps:
_, _, motion_pulse_steps, motion_pulse_noise = _prepare_dynamic_sigmas(
sigmas, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct,
sched, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct,
motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength)
motion_pulse_noise = max(0.0, float(motion_pulse_noise or 0.0))
motion_step_noise = {
@@ -1150,14 +1174,13 @@ def _sample_const_rf_full(model, x, sigmas, extra_args, callback, disable,
for item in (motion_pulse_steps or []) if isinstance(item, dict)
}
late_start = _get_late_start_index(total_steps, high_quality_pct)
quality_sigma_start = float(sigmas[late_start].item()) if late_start < sigmas.shape[0] else None
num_quality_steps = total_steps - late_start
late_start = _get_late_start_index(sched_steps, high_quality_pct)
quality_sigma_start = float(sched[late_start].item()) if late_start < sched.shape[0] else None
num_quality_steps = sched_steps - late_start
s_in = x.new_ones([x.shape[0]])
prev_denoised = None
prev_h = None
quality_step_index = 0
_kvlock_state = {}
# Audio-safe sampling: on a packed LTXAV latent, keep ancestral noise + steering on the
@@ -1178,11 +1201,12 @@ def _sample_const_rf_full(model, x, sigmas, extra_args, callback, disable,
for i in comfy.utils.model_trange(total_steps, disable=disable):
sigma = sigmas[i]
sigma_next = sigmas[i + 1]
step_abs = i + step_offset # position in the REAL schedule, not in this call's slice
in_quality = quality_sigma_start is not None and float(sigma.item()) <= quality_sigma_start
velocity_target = _velocity_bias_target(sigmas, sigma)
velocity_target = _velocity_bias_target(sched, sigma)
if not in_quality:
pulse = motion_step_noise.get(int(i), 0.0)
pulse = motion_step_noise.get(int(step_abs), 0.0)
if pulse > 0.0:
x = _video_only(_apply_motion_pulse(x, sigma, sigma_next, pulse, noise_sampler), x, video_mask)
prev_denoised = None
@@ -1191,7 +1215,7 @@ def _sample_const_rf_full(model, x, sigmas, extra_args, callback, disable,
if _velocity_bias_enabled(velocity_bias_mode, "apply"):
x_pre = x
x = _apply_velocity_bias(x, velocity_refinement_key, velocity_target, velocity_bias_strength,
sigma_ratio=_sigma_ratio(sigmas, sigma),
sigma_ratio=_sigma_ratio(sched, sigma),
prompt_sig=rescue_prompt_sig, source=velocity_bias_source)
x = _video_only(x, x_pre, video_mask)
@@ -1245,7 +1269,10 @@ def _sample_const_rf_full(model, x, sigmas, extra_args, callback, disable,
if num_quality_steps <= 1:
effective_blend = correction_blend
else:
effective_blend = 0.0 if quality_step_index < (num_quality_steps // 2) else correction_blend
# Which quality step this is IN THE SCHEDULE — a per-call counter would
# restart the euler->Heun ramp inside each half of an anchor_shift split.
q_idx = max(0, step_abs - late_start)
effective_blend = 0.0 if q_idx < (num_quality_steps // 2) else correction_blend
dt = sigma_next - sigma
d1 = k_diffusion_sampling.to_d(x, sigma, denoised_eff)
if effective_blend > 0.0:
@@ -1261,11 +1288,10 @@ def _sample_const_rf_full(model, x, sigmas, extra_args, callback, disable,
# Heun changed x with a corrected direction; invalidate AB2 history.
prev_denoised = None
prev_h = None
quality_step_index += 1
else:
# Early phase: ancestral RF euler (matches sample_euler_ancestral_RF), with
# eta decay toward the quality boundary, using the AB2 denoised estimate.
effective_eta = _effective_eta(eta, eta_final, sigmas, sigma)
effective_eta = _effective_eta(eta, eta_final, sched, sigma)
# Deterministic euler-RF step to sigma_next (this is what audio rides — no
# ancestral re-noising). For single-stream video (mask None) it is only used
# as the eta==0 fallback; the full ancestral result is kept for video.
@@ -1310,7 +1336,8 @@ def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=N
rescue_prompt_sig=None,
eta_final=1.0,
normalize_strength=0.0,
normalize_start_sigma=0.9):
normalize_start_sigma=0.9,
schedule_context=None):
"""
Hybrid sampler:
- Early schedule: Euler ancestral with order-2 denoised extrapolation for
@@ -1387,6 +1414,7 @@ def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=N
velocity_bias_source=velocity_bias_source,
normalize_strength=normalize_strength,
normalize_start_sigma=normalize_start_sigma,
schedule_context=schedule_context,
)
extra_args = {} if extra_args is None else extra_args
@@ -1758,7 +1786,7 @@ def sample_funpack_distilled_flow(model, x, sigmas, extra_args=None, callback=No
alg_guide_tail_frames=0,
alg_guide_blur_strength=2.0, alg_guide_blur_sigma_threshold=0.975,
mg_enabled=False, mg_strength=0.5, mg_decay=0.5, mg_sigma_threshold=0.975,
quality_sharpness=0.0):
quality_sharpness=0.0, schedule_context=None):
"""
ODE sampler for distilled few-step video models (e.g. LTX2.3 distilled LoRA).
@@ -1808,10 +1836,18 @@ def sample_funpack_distilled_flow(model, x, sigmas, extra_args=None, callback=No
if total_steps <= 0:
return x
# `sigmas` is not always the schedule the user configured: anchor_shift runs a scene as
# two sample_custom calls, so this can be a FRAGMENT. Everything phase-relative below —
# where the final-correction window opens, how far along the AB2 ramp is, what the
# velocity-bias sigma ratio is measured against — has to be read off the REAL schedule,
# or enabling anchor_shift silently re-times all of it inside each half.
sched, step_offset = _schedule_view(sigmas, schedule_context)
sched_steps = max(1, int(len(sched)) - 1)
order = max(1, min(2, int(order)))
s_noise = max(0.0, min(0.5, float(s_noise)))
final_correction_steps = max(0, min(total_steps // 2, int(final_correction_steps)))
correction_start_idx = total_steps - final_correction_steps
final_correction_steps = max(0, min(sched_steps // 2, int(final_correction_steps)))
correction_start_idx = sched_steps - final_correction_steps
_RESCUE_LOG["warned_no_memory"] = False
_RESCUE_LOG["warned_no_prompt_match"] = False
@@ -1918,7 +1954,9 @@ def sample_funpack_distilled_flow(model, x, sigmas, extra_args=None, callback=No
for i in comfy.utils.model_trange(total_steps, disable=disable):
sigma = sigmas[i]
sigma_next = sigmas[i + 1]
velocity_target = _velocity_bias_target(sigmas, sigma)
# Position in the REAL schedule, not in this call's slice of it.
step_abs = i + step_offset
velocity_target = _velocity_bias_target(sched, sigma)
if alg_active:
model.latent_image = alg_latents[(
@@ -1929,7 +1967,7 @@ def sample_funpack_distilled_flow(model, x, sigmas, extra_args=None, callback=No
if _velocity_bias_enabled(velocity_bias_mode, "apply"):
x_pre = x
x = _apply_velocity_bias(x, velocity_refinement_key, velocity_target, velocity_bias_strength,
sigma_ratio=_sigma_ratio(sigmas, sigma),
sigma_ratio=_sigma_ratio(sched, sigma),
prompt_sig=rescue_prompt_sig, source=velocity_bias_source)
x = _video_only(x, x_pre, video_mask)
@@ -1947,7 +1985,7 @@ def sample_funpack_distilled_flow(model, x, sigmas, extra_args=None, callback=No
# Restore high-frequency detail during the final Heun-correction steps (free,
# video-only) — same unsharp mechanism as the Hybrid Euler 2S sampler.
if i >= correction_start_idx:
if step_abs >= correction_start_idx:
denoised = _video_only(_apply_quality_sharpness(denoised, prev_denoised, quality_sharpness), denoised, video_mask)
# Video-only latent normalization (opt-in, anti-overbake) — stacks on top of the ODE.
@@ -1976,8 +2014,8 @@ def sample_funpack_distilled_flow(model, x, sigmas, extra_args=None, callback=No
# the schedule. Early/high-sigma steps stay near 1st-order euler (the denoised estimate
# is rough there and full AB2 overshoots); late/detail steps get full AB2. Reuses the
# already-computed AB2 estimate, so no extra model evals. No-op at order=1.
if ab2_ramp and total_steps > 1:
w = i / (total_steps - 1) # 0 at first step -> 1 at last
if ab2_ramp and sched_steps > 1:
w = step_abs / (sched_steps - 1) # 0 at first step -> 1 at last
denoised_eff = denoised + (denoised_eff - denoised) * w
# Audio rides plain 1st-order euler: keep the (ramped) AB2 estimate for video,
@@ -1995,7 +2033,7 @@ def sample_funpack_distilled_flow(model, x, sigmas, extra_args=None, callback=No
dt = sigma_next - sigma # negative: sigmas decrease
if i >= correction_start_idx:
if step_abs >= correction_start_idx:
# Heun predictor-corrector.
# Predictor: Euler step using the (multistep-corrected) denoised.
d1 = k_diffusion_sampling.to_d(x, sigma, denoised_eff)
@@ -2951,7 +2989,8 @@ class FunPackLTXAVSceneChainSampler:
def _sample_chunk(self, model, sampler, sigmas, seed, cfg, positive, negative, latent,
pbar=None, step_offset=0, alg_guide_tail_frames=0,
alg_guide_blur_strength=2.0, alg_guide_blur_sigma_threshold=0.975,
bounded_attention_enabled=False, continue_from_state=False):
bounded_attention_enabled=False, continue_from_state=False,
schedule_context=None):
if sampler is None:
raise ValueError("sampler input is required.")
if not isinstance(sigmas, torch.Tensor):
@@ -2982,6 +3021,13 @@ class FunPackLTXAVSceneChainSampler:
extra_options["alg_guide_tail_frames"] = int(alg_guide_tail_frames)
extra_options["alg_guide_blur_strength"] = float(alg_guide_blur_strength)
extra_options["alg_guide_blur_sigma_threshold"] = float(alg_guide_blur_sigma_threshold)
# anchor_shift runs a scene as two calls over slices of one schedule. Tell the sampler
# what the real schedule is (and where this slice sits in it) so its phase boundaries
# stay where the user's schedule put them instead of being re-derived per slice — see
# _schedule_view. Always written, even as None, so it can never leak into a later chunk.
if isinstance(extra_options, dict) and sampler.sampler_function in (
sample_funpack_distilled_flow, sample_funpack_hybrid_euler_2s):
extra_options["schedule_context"] = schedule_context
# EXPERIMENTAL Bounded Attention: model-level attention hooks (sampler-agnostic, unlike
# the toggles above which only work on Distilled Flow), so install/remove here rather
@@ -3057,6 +3103,16 @@ class FunPackLTXAVSceneChainSampler:
Both passes need at least one real step, so a threshold outside the schedule's
range is refused rather than silently producing a zero-step pass. A restart BELOW
the cut is refused too — that would mean removing noise, which is not a thing.
Everything here is read off the ACTUAL schedule, never off the typed floats: walk
the sigmas and the last pre-shift step is the one whose NEXT sigma has crossed the
threshold. On a 1.000, 0.955, 0.893, 0.812, 0.715, 0.603, 0.482, 0.241, 0.121, 0.0
schedule with shift 0.482, steps at 1.000 and 0.955 both have a next sigma still
above it and proceed; the step at 0.603 has next=0.482, so it is the last one and
the shift happens after it. `_first_at_or_below` is that walk in closed form — the
first index at or below the threshold IS the sigma the last pre-shift step lands on.
Returned cut/restart values are the schedule's own float32 entries, so the sigma
pass 2 re-enters at is exactly the one _anchor_shift_renoise rebuilds the state for.
"""
if not isinstance(sigmas, torch.Tensor) or sigmas.numel() < 3:
return None
@@ -3087,7 +3143,9 @@ class FunPackLTXAVSceneChainSampler:
first, second, resume = sigmas[:cut + 1], sigmas[start:], start == cut
if first.numel() < 2 or second.numel() < 2:
return None
return first, second, vals[cut], vals[start], resume
# `start` goes back to the caller too: it is pass 2's offset into the real schedule,
# which the sampler needs to keep its phase boundaries where the schedule put them.
return first, second, vals[cut], vals[start], resume, start
def _anchor_shift_renoise(self, latent, from_sigma, to_sigma, seed, reset_indices=()):
"""Move a latent that sits mid-trajectory at `from_sigma` back up to `to_sigma`.
@@ -5534,10 +5592,14 @@ class FunPackLTXAVSceneChainSampler:
if _shift_split is None:
run_mechanisms.append(f"anchor_shift(SKIPPED: {_reason})")
if _shift_split is not None:
_sig_a, _sig_b, _cut_at, _restart_at, _resume = _shift_split
_sig_a, _sig_b, _cut_at, _restart_at, _resume, _restart_idx = _shift_split
# Both passes are slices of THIS schedule. Hand the sampler the real thing
# plus where each slice starts, so the quality phase, final-correction
# window, AB2 ramp, pulse anchors and eta decay stay where the schedule
# puts them instead of being re-derived inside each half.
_pass1 = self._sample_chunk(
model, sampler, _sig_a, scene_seed, cfg, scene_positive, scene_negative,
chunk, **_sample_kwargs)
chunk, schedule_context={"sigmas": sigmas, "offset": 0}, **_sample_kwargs)
_shifted, _dropped = self._anchor_shift_latent(
_pass1, _shift_drop, anchor_shift_tail, anchor_shift_fresh_audio)
if not _resume:
@@ -5556,6 +5618,7 @@ class FunPackLTXAVSceneChainSampler:
_kw2 = dict(_sample_kwargs)
_kw2["step_offset"] = _sample_kwargs["step_offset"] + (_sig_a.numel() - 1)
_kw2["continue_from_state"] = True
_kw2["schedule_context"] = {"sigmas": sigmas, "offset": _restart_idx}
sampled = self._sample_chunk(
model, sampler, _sig_b, scene_seed + 4242, cfg, scene_positive,
scene_negative, _shifted, **_kw2)
tests/test_anchor_shift.py
+51 -3
View File
@@ -55,7 +55,7 @@ def _latent(frames=10, audio=True):
# ── schedule splitting ──────────────────────────────────────────────────────
def test_continue_mode_stops_at_the_cut_and_resumes_from_that_state():
first, second, cut_at, restart_at, resume = _node()._anchor_shift_split_sigmas(
first, second, cut_at, restart_at, resume, start_idx = _node()._anchor_shift_split_sigmas(
SCHEDULE, 0.909, 0.0)
# float32: the schedule stores 0.909 as 0.90899997, so compare loosely.
assert abs(cut_at - 0.909) < 1e-5 and abs(restart_at - 0.909) < 1e-5
@@ -72,7 +72,7 @@ def test_continue_mode_stops_at_the_cut_and_resumes_from_that_state():
def test_rewind_mode_still_stops_pass_one_at_the_cut():
"""The cut sigma is the whole feature — rewind must NOT override it. Pass 1 stops
where the user asked, and the restart only decides where pass 2 re-enters."""
first, second, cut_at, restart_at, resume = _node()._anchor_shift_split_sigmas(
first, second, cut_at, restart_at, resume, start_idx = _node()._anchor_shift_split_sigmas(
SCHEDULE, 0.725, 0.975)
assert abs(cut_at - 0.725) < 1e-5 and abs(restart_at - 0.975) < 1e-5
assert [round(v, 3) for v in first.tolist()] == [1.0, 0.975, 0.909, 0.725]
@@ -89,12 +89,60 @@ def test_rewind_refuses_a_restart_below_the_cut():
def test_restart_landing_on_the_cut_is_just_continue():
"""A restart that snaps to the cut's own step has nothing to re-noise."""
_, _, cut_at, restart_at, resume = _node()._anchor_shift_split_sigmas(
_, _, cut_at, restart_at, resume, _ = _node()._anchor_shift_split_sigmas(
SCHEDULE, 0.909, 0.909)
assert abs(cut_at - restart_at) < 1e-5
assert resume is True
# ── the sampler must still see the REAL schedule ────────────────────────────
# The user's own 9-step schedule, used here because the boundary walk has to be right on
# the schedule they actually run, not on the toy one above.
USER_SCHEDULE = torch.tensor(
[1.0, 0.955, 0.893, 0.812, 0.715, 0.603, 0.482, 0.241, 0.121, 0.0])
def test_cut_is_the_step_whose_next_sigma_crosses_the_threshold():
"""Walk the real schedule: the last pre-shift step is the one whose NEXT sigma has
crossed the shift point. At shift 0.482 the steps at 1.000 and 0.955 both have a next
sigma still above it and proceed; the step at 0.603 has next=0.482, so it is the last."""
first, second, cut_at, restart_at, resume, start_idx = _node()._anchor_shift_split_sigmas(
USER_SCHEDULE, 0.482, 0.812)
assert [round(v, 3) for v in first.tolist()] == [1.0, 0.955, 0.893, 0.812, 0.715, 0.603, 0.482]
assert abs(cut_at - 0.482) < 1e-5
# Pass 2 re-enters at 0.812, which is index 3 of the real schedule.
assert start_idx == 3
assert abs(restart_at - 0.812) < 1e-5
assert [round(v, 3) for v in second.tolist()] == [0.812, 0.715, 0.603, 0.482, 0.241, 0.121, 0.0]
assert resume is False
def test_schedule_view_hands_the_sampler_the_whole_schedule():
"""Without this the sampler measures its phases from the SLICE: the final-correction
window reopens at the end of pass 1, the AB2 ramp restarts, and the velocity-bias
ratio is taken against the slice's own first sigma instead of the schedule's."""
slice_ = USER_SCHEDULE[3:]
sched, offset = samplers._schedule_view(
slice_, {"sigmas": USER_SCHEDULE, "offset": 3})
assert sched is USER_SCHEDULE and offset == 3
# Concretely: 9 real steps, not the 6 the slice would have claimed.
assert len(sched) - 1 == 9 and len(slice_) - 1 == 6
def test_schedule_view_falls_back_to_the_sigmas_it_was_given():
"""Every normal single-call run passes no context — `sigmas` IS the whole schedule."""
for ctx in (None, {}, {"sigmas": None}, {"sigmas": torch.tensor([1.0])}, "nonsense"):
sched, offset = samplers._schedule_view(USER_SCHEDULE, ctx)
assert sched is USER_SCHEDULE and offset == 0
def test_schedule_view_never_returns_a_negative_offset():
sched, offset = samplers._schedule_view(
USER_SCHEDULE, {"sigmas": USER_SCHEDULE, "offset": -5})
assert offset == 0
# ── the mid-trajectory re-noise ─────────────────────────────────────────────
class _FakeNested(list):