import copy import hashlib import math import os import time as _time import types import torch import comfy.k_diffusion.sampling as k_diffusion_sampling import comfy.model_sampling import comfy.nested_tensor import comfy.sample import comfy.samplers import comfy.utils try: from . import funpack_log as _log except ImportError: # flat import when ComfyUI loads the pack as a top-level module import funpack_log as _log MOTION_PULSE_MODES = ["off", "balanced", "aggressive", "custom"] VELOCITY_BIAS_MODES = ["off", "capture", "apply", "capture_and_apply"] # Normalized-sigma targets (current/sigmas[0]) where velocity is captured and rescue acts. # Chosen to land on the structure AND detail-forming steps of the standard 8-step LTX # schedule [..,0.909,0.725,0.422,0] (matched within +/-0.065): 0.90->0.909 (structure), # 0.72->0.725 and 0.42->0.422 (where runs actually go good/bad). The old (0.90, 0.80) only # ever matched 0.909, so rescue fired once per gen at a near-noise step. Velocity-bias APPLY # still fires only at the highest target (structure); the extra targets are for rescue. VELOCITY_BIAS_TARGETS = (0.90, 0.72, 0.42) # Good-trajectory bank (rating-promoted). Also read by velocity-bias "apply". VELOCITY_BIAS_MEMORY = {} # Bad-trajectory bank (promoted from Awful/dislike), used by rescue to steer away. VELOCITY_BIAS_BAD_MEMORY = {} # Per-key staging of the most recently captured trajectory, awaiting its rating. # capture stages here; the refiner commits it to the good or bad bank (or drops it) # once the user rates that generation — so neither bank is rating-blind. VELOCITY_BIAS_PENDING = {} # Refinement keys whose banks have been loaded from disk this process (load-once). _VELOCITY_LOADED = set() def _velocity_store_path(refinement_key): """On-disk store for a key's velocity/rescue banks, alongside the other per-key refinement data. Persisted so the learned good/bad trajectories survive restarts.""" base = os.path.dirname(os.path.abspath(__file__)) d = os.path.join(base, "refinements", "attn_maps") os.makedirs(d, exist_ok=True) norm = str(refinement_key or "default").strip() or "default" h = hashlib.md5(f"velocity::{norm}".encode()).hexdigest()[:16] return os.path.join(d, f"velocity_{h}.pt") def _serialize_velocity_bank(bank, norm): """weights_only-safe representation (lists/dicts/str/int + tensors; no tuple keys).""" out = [] for key, slot in bank.items(): if not (isinstance(key, tuple) and len(key) == 3 and key[0] == norm): continue direction = slot.get("direction") out.append({ "refkey": str(key[0]), "target": str(key[1]), "shape": [int(v) for v in key[2]], "count": int(slot.get("count", 0)), "direction": direction.half() if isinstance(direction, torch.Tensor) else None, "clusters": [ {"sig": c["sig"].half() if isinstance(c.get("sig"), torch.Tensor) else None, "direction": c["direction"].half() if isinstance(c.get("direction"), torch.Tensor) else None, "count": int(c.get("count", 0))} for c in slot.get("clusters", []) if isinstance(c, dict) ], }) return out def _deserialize_velocity_into(bank, entries): for e in entries or []: try: # Old stores included an "aspect" field — ignored now (key dropped it). key = (str(e["refkey"]), str(e["target"]), tuple(int(v) for v in e["shape"])) d = e.get("direction") bank[key] = { "count": int(e.get("count", 0)), "direction": d.float() if isinstance(d, torch.Tensor) else None, "clusters": [ {"sig": c["sig"].float() if isinstance(c.get("sig"), torch.Tensor) else None, "direction": c["direction"].float() if isinstance(c.get("direction"), torch.Tensor) else None, "count": int(c.get("count", 0))} for c in e.get("clusters", []) if isinstance(c, dict) ], } except Exception: continue def _ensure_velocity_loaded(refinement_key): """Load a key's persisted banks into memory once per process.""" norm = str(refinement_key or "default").strip() or "default" if norm in _VELOCITY_LOADED: return _VELOCITY_LOADED.add(norm) path = _velocity_store_path(norm) if not os.path.exists(path): return try: data = torch.load(path, map_location="cpu", weights_only=True) _deserialize_velocity_into(VELOCITY_BIAS_MEMORY, data.get("good", [])) _deserialize_velocity_into(VELOCITY_BIAS_BAD_MEMORY, data.get("bad", [])) ng = sum(1 for k in VELOCITY_BIAS_MEMORY if isinstance(k, tuple) and k[0] == norm) nb = sum(1 for k in VELOCITY_BIAS_BAD_MEMORY if isinstance(k, tuple) and k[0] == norm) print(f"[FunPack rescue] loaded persisted velocity memory for '{norm}': {ng} good, {nb} bad bucket(s)") except Exception as e: print(f"[FunPack rescue] load failed: {e}") def _save_velocity_store(refinement_key): """Persist a key's good+bad banks to disk (called after a commit changes them).""" norm = str(refinement_key or "default").strip() or "default" path = _velocity_store_path(norm) try: good = _serialize_velocity_bank(VELOCITY_BIAS_MEMORY, norm) bad = _serialize_velocity_bank(VELOCITY_BIAS_BAD_MEMORY, norm) if not good and not bad: if os.path.exists(path): os.remove(path) return torch.save({"good": good, "bad": bad}, path) except Exception as e: print(f"[FunPack rescue] save failed: {e}") # Prompt-conditioned rescue clustering. The good-trajectory memory keeps, per # velocity key, a set of clusters each tagged with a pooled+normalized prompt # signature. Capture merges into the nearest cluster (or starts a new one); # rescue builds its reference from clusters whose prompt is similar enough to the # current run, so it steers toward what was good *for this kind of prompt* rather # than a prompt-blind global average. RESCUE_CLUSTER_MERGE_SIM = 0.92 # capture: merge into a cluster at/above this cosine RESCUE_CLUSTER_MAX = 8 # cap clusters per key (evict lowest-count); each holds a full-shape direction RESCUE_USE_SIM = 0.5 # rescue: only use clusters at/above this cosine RESCUE_MIN_COUNT = 2 # rescue: need this many matched captures to act def _normalize_sig(sig): """Pooled prompt vector -> CPU float unit vector, or None.""" if not isinstance(sig, torch.Tensor) or sig.numel() == 0: return None s = sig.detach().float().reshape(-1).cpu() n = s.norm() if not torch.isfinite(n) or float(n) <= 1e-8: return None return s / n def _update_prompt_clusters(slot, prompt_sig, direction): """Capture-side: fold (prompt_sig, direction) into the slot's prompt clusters.""" if prompt_sig is None: return clusters = slot.setdefault("clusters", []) best_i, best_cos = -1, -1.0 for i, c in enumerate(clusters): cs = c.get("sig") if not isinstance(cs, torch.Tensor) or cs.shape != prompt_sig.shape: continue cos = float(cs @ prompt_sig) if cos > best_cos: best_cos, best_i = cos, i if best_i >= 0 and best_cos >= RESCUE_CLUSTER_MERGE_SIM: c = clusters[best_i] n = int(c.get("count", 1)) if isinstance(c.get("direction"), torch.Tensor) and c["direction"].shape == direction.shape: c["direction"] = (c["direction"] * n + direction) / float(n + 1) merged = c["sig"] * n + prompt_sig mn = merged.norm().clamp_min(1e-8) c["sig"] = merged / mn c["count"] = min(n + 1, 256) else: clusters.append({"sig": prompt_sig.clone(), "direction": direction.clone(), "count": 1}) if len(clusters) > RESCUE_CLUSTER_MAX: clusters.sort(key=lambda c: int(c.get("count", 0)), reverse=True) del clusters[RESCUE_CLUSTER_MAX:] def _select_prompt_direction(slot, prompt_sig, x_shape): """Rescue-side: similarity-weighted good direction over prompt-relevant clusters. Returns None when there is no confident prompt-matching reference (caller then skips rescue rather than steering toward a prompt-blind average).""" clusters = slot.get("clusters") if isinstance(slot, dict) else None if not clusters or prompt_sig is None: return None num = None total_w = 0.0 matched = 0 for c in clusters: cs = c.get("sig") d = c.get("direction") if not isinstance(cs, torch.Tensor) or cs.shape != prompt_sig.shape: continue if not isinstance(d, torch.Tensor) or tuple(d.shape) != tuple(x_shape): continue cos = float(cs @ prompt_sig) if cos < RESCUE_USE_SIM: continue contrib = d * cos num = contrib if num is None else num + contrib total_w += cos matched += int(c.get("count", 1)) if num is None or total_w <= 0.0 or matched < RESCUE_MIN_COUNT: return None return num / total_w def _nearest_cluster_direction(slot, prompt_sig, x_shape): """Option D: the single highest-similarity prompt cluster's direction (>= RESCUE_USE_SIM). Returns one cluster's direction rather than a blend across clusters, preserving that real good trajectory's high-frequency content instead of averaging several together (averaging is what low-passes / softens). None if no cluster matches confidently. """ clusters = slot.get("clusters") if isinstance(slot, dict) else None if not clusters or prompt_sig is None: return None best_d, best_cos = None, RESCUE_USE_SIM for c in clusters: cs = c.get("sig") d = c.get("direction") if not isinstance(cs, torch.Tensor) or cs.shape != prompt_sig.shape: continue if not isinstance(d, torch.Tensor) or tuple(d.shape) != tuple(x_shape): continue cos = float(cs @ prompt_sig) if cos >= best_cos: best_cos, best_d = cos, d return best_d def _velocity_direction(slot, prompt_sig, x_shape, source="mean"): """Resolve a steering direction from a bank slot. source="nearest" (option D): single best-matching prompt cluster — preserves one real good trajectory's detail. Falls back to the global mean when no cluster matches (e.g. no prompt sig, or memory captured before clustering existed). source="mean" (default): prompt-blind global average (legacy behavior). """ if source == "nearest": d = _nearest_cluster_direction(slot, prompt_sig, x_shape) if d is not None: return d g = slot.get("direction") if isinstance(slot, dict) else None return g if isinstance(g, torch.Tensor) and tuple(g.shape) == tuple(x_shape) else None def _sigma_fn(t): return t.neg().exp() def _t_fn(sigma): return sigma.log().neg() def _hybrid_ode_step(model, x, sigma, sigma_next, s_in, extra_args, correction_blend, denoised=None): if denoised is None: denoised = model(x, sigma * s_in, **extra_args) if sigma_next == 0: return denoised, denoised d = k_diffusion_sampling.to_d(x, sigma, denoised) dt = sigma_next - sigma x_euler = x + d * dt if correction_blend <= 0.0: return x_euler, denoised t = _t_fn(sigma) t_next = _t_fn(sigma_next) h = t_next - t r = 0.5 s = t + r * h x_mid = (_sigma_fn(s) / _sigma_fn(t)) * x - torch.expm1(-h * r) * denoised denoised_mid = model(x_mid, _sigma_fn(s) * s_in, **extra_args) x_2s = (_sigma_fn(t_next) / _sigma_fn(t)) * x - torch.expm1(-h) * denoised_mid if correction_blend >= 1.0: return x_2s, denoised return x_euler.lerp(x_2s, correction_blend), denoised def _apply_motion_pulse(x, sigma, sigma_next, pulse_noise, noise_sampler): if pulse_noise <= 0.0 or sigma_next >= sigma: return x sigma_delta_sq = max(0.0, float(sigma * sigma - sigma_next * sigma_next)) if sigma_delta_sq <= 0.0: return x sigma_delta = math.sqrt(sigma_delta_sq) return x + noise_sampler(sigma, sigma_next) * (pulse_noise * sigma_delta) def clear_velocity_bias_memory(refinement_key): """Drop all stored velocity-bias / rescue trajectory memory (global averages and prompt clusters) for a key. Called on Studio session reset so rescue and velocity bias start from a clean slate instead of steering toward last session's content.""" norm = str(refinement_key or "default").strip() or "default" # Make sure persisted buckets are in memory first, so the on-disk store is rewritten # empty (deleted) rather than left behind. _ensure_velocity_loaded(norm) removed = 0 for bank in (VELOCITY_BIAS_MEMORY, VELOCITY_BIAS_BAD_MEMORY): for k in [k for k in list(bank) if isinstance(k, tuple) and k and k[0] == norm]: bank.pop(k, None) removed += 1 VELOCITY_BIAS_PENDING.pop(norm, None) # Delete the on-disk store too (this is the deliberate Session-reset path). try: path = _velocity_store_path(norm) if os.path.exists(path): os.remove(path) except Exception as e: print(f"[FunPack rescue] disk clear failed: {e}") if removed: print(f"[FunPack rescue] cleared {removed} velocity-bias memory bucket(s) " f"(good+bad) for key '{norm}' (session reset)") return removed def _velocity_bias_enabled(mode, action): mode = (mode or "off").lower() if mode == "capture_and_apply": return action in {"capture", "apply"} return mode == action def _velocity_bias_key(refinement_key, target, x): # (refinement_key, target, latent_shape). Shape already encodes resolution, so a # separate aspect bucket was redundant — dropped. shape = tuple(int(item) for item in getattr(x, "shape", ())) key = str(refinement_key or "default").strip() or "default" return (key, f"{float(target):.2f}", shape) def _sigma_ratio(sigmas, sigma): try: start = float(sigmas[0].item()) current = float(sigma.item()) except Exception: return None if start <= 0.0: return None return current / start def _effective_eta(eta, eta_final, sigmas, sigma): """Ancestral eta decayed across the schedule: full ``eta`` at the start (sigma == sigmas[0]) ramping linearly to ``eta_final`` near the end (sigma -> 0). Anchored to raw schedule progress (sigma/sigmas[0]), NOT the quality-phase boundary. The old anchor collapsed to sigma=0 when high_quality_pct=0, silently pinning effective_eta at ``eta`` for every step (eta_final ignored) and leaving full-strength ancestral noise on the final detail steps with no deterministic phase to resolve it. """ if eta_final >= eta: return eta ratio = _sigma_ratio(sigmas, sigma) if ratio is None: return eta ratio = max(0.0, min(1.0, ratio)) return eta_final + (eta - eta_final) * ratio def _velocity_bias_target(sigmas, sigma): ratio = _sigma_ratio(sigmas, sigma) if ratio is None: return None target = min(VELOCITY_BIAS_TARGETS, key=lambda item: abs(float(item) - ratio)) return target if abs(float(target) - ratio) <= 0.065 else None def _capture_velocity_bias(refinement_key, target, x, sigma, denoised, prompt_sig=None): """Stage this step's trajectory for the current generation. It is NOT committed to any bank here — the refiner commits it to good or bad once the gen is rated, so the banks never absorb unrated (possibly awful) runs. Staging overwrites per (key, shape, target) so only the latest generation for a key is pending.""" if target is None: return try: direction = k_diffusion_sampling.to_d(x, sigma, denoised).detach().float().cpu() except Exception: return key = _velocity_bias_key(refinement_key, target, x) norm = key[0] staged = VELOCITY_BIAS_PENDING.setdefault(norm, {}) staged[key] = {"direction": direction, "sig": prompt_sig.clone() if isinstance(prompt_sig, torch.Tensor) else None} def _merge_into_bank(bank, key, direction, prompt_sig): """Fold one staged trajectory into a bank's global average + prompt clusters.""" slot = bank.setdefault(key, {"count": 0, "direction": None}) count = int(slot.get("count", 0)) previous = slot.get("direction") if count <= 0 or not isinstance(previous, torch.Tensor) or tuple(previous.shape) != tuple(direction.shape): slot["direction"] = direction slot["count"] = 1 else: slot["direction"] = (previous * count + direction) / float(count + 1) slot["count"] = min(count + 1, 256) _update_prompt_clusters(slot, prompt_sig, direction) def commit_staged_velocity(refinement_key, verdict): """Commit the staged trajectory for a key to the good or bad bank based on the generation's rating. verdict: 'good' | 'bad' | 'drop'. Called by the refiner when it processes the previous gen's rating. Always clears the staging slot.""" norm = str(refinement_key or "default").strip() or "default" staged = VELOCITY_BIAS_PENDING.pop(norm, None) if not staged or verdict == "drop": return 0 _ensure_velocity_loaded(norm) # merge on top of any persisted buckets bank = VELOCITY_BIAS_MEMORY if verdict == "good" else VELOCITY_BIAS_BAD_MEMORY n = 0 for key, entry in staged.items(): d = entry.get("direction") if isinstance(d, torch.Tensor): _merge_into_bank(bank, key, d, entry.get("sig")) n += 1 if n: _save_velocity_store(norm) # persist so it survives restarts print(f"[FunPack rescue] committed {n} staged trajectory bucket(s) to " f"'{verdict}' bank for key '{norm}' (persisted)") return n def _apply_velocity_bias(x, refinement_key, target, strength, sigma_ratio=None, prompt_sig=None, source="mean"): """Steer the latent toward the averaged good-trajectory direction. Two anti-softening guards (see velocity-bias quality findings): C. Structure-only + sigma decay — only fires at the highest velocity target (structure-forming sigma, ~0.90 ratio); the lower detail target is left untouched so high-frequency detail forms freely. Within that, strength decays with the sigma ratio so the push is strongest early and fades as sigma drops. B. Magnitude-preserving — instead of adding the delta as raw energy (x + delta, which pulls the latent toward the low-variance mean and progressively softens), we rotate x toward (x + delta) and renormalize back to |x|. Direction is steered; no energy is injected. """ if target is None or strength <= 0.0: return x # C: structure target only — skip the lower (detail) targets entirely. if float(target) < max(VELOCITY_BIAS_TARGETS) - 1e-6: return x _ensure_velocity_loaded(refinement_key) key = _velocity_bias_key(refinement_key, target, x) slot = VELOCITY_BIAS_MEMORY.get(key) if not isinstance(slot, dict): return x # D: nearest single good trajectory (prompt cluster) vs prompt-blind global mean. direction = _velocity_direction(slot, prompt_sig, x.shape, source=source) if not isinstance(direction, torch.Tensor): return x try: # C: decay strength with the sigma ratio (full at the structure target, less below). decay = 1.0 if sigma_ratio is not None: decay = max(0.0, min(1.0, float(sigma_ratio) / max(VELOCITY_BIAS_TARGETS))) eff_strength = max(0.0, min(3.0, float(strength))) * decay if eff_strength <= 0.0: return x direction = direction.to(device=x.device, dtype=x.dtype) delta = direction * eff_strength x_norm = x.detach().float().norm().clamp_min(1e-8) # Cap the rotation toward the remembered direction, scaled by strength so the WHOLE # slider is meaningful (the old fixed 4.5% saturated almost immediately, making the # upper range a no-op). 0.30*strength keeps the low end ~unchanged (strength~0.15 == # old ~4.5%) while high strength approaches full action replacement — capped at 0.95 # so it never totally wipes the current gen. Magnitude-preserving renorm keeps |x|. max_delta = x_norm * min(0.95, 0.30 * eff_strength) delta_norm = delta.detach().float().norm().clamp_min(1e-8) if delta_norm > max_delta: delta = delta * (max_delta / delta_norm).to(device=x.device, dtype=x.dtype) # B: rotate, don't add — renormalize the biased latent back to the original norm. biased = x + delta biased_norm = biased.detach().float().norm().clamp_min(1e-8) return biased * (x_norm / biased_norm).to(device=x.device, dtype=x.dtype) except Exception as _e: _log.failed("FunPackSceneChain", "velocity bias", _e, "this step runs unbiased — motion will read as it would with the knob at 0") return x # Per-run rescue logging state, reset at the start of each sampler call. _RESCUE_LOG = {"warned_no_memory": False, "warned_no_prompt_match": False, "fired": 0} def _rescue_reference(bank, key, prompt_sig, x_shape, source="mean"): """Resolve a steering direction from a bank. source="nearest" (option D): single best-matching prompt cluster (one real good trajectory's detail), falling back to the global average if none matches. source="mean" (default): similarity-weighted blend over prompt-relevant clusters when a prompt signature is available, else the prompt-blind global average. """ slot = bank.get(key) if not isinstance(slot, dict): return None if source == "nearest": return _velocity_direction(slot, prompt_sig, x_shape, source="nearest") if prompt_sig is not None: return _select_prompt_direction(slot, prompt_sig, x_shape) g = slot.get("direction") return g if isinstance(g, torch.Tensor) and tuple(g.shape) == tuple(x_shape) else None def _rescue_denoised(denoised, x, sigma, refinement_key, target, threshold, strength, prompt_sig=None, source="mean"): """Reactive in-flight rescue, rating-aware. Pulls this step's trajectory toward the GOOD bank (rating-promoted likes) and away from the BAD bank (Awful/dislike). Fires when the trajectory has diverged from good beyond ``threshold`` OR aligned with bad beyond ``threshold``. Magnitude preserved (no energy injected). No matching reference in either bank -> no-op. """ if target is None or strength <= 0.0: return denoised sig = float(sigma.item()) if isinstance(sigma, torch.Tensor) else float(sigma) if sig <= 1e-6: return denoised _ensure_velocity_loaded(refinement_key) key = _velocity_bias_key(refinement_key, target, x) good = _rescue_reference(VELOCITY_BIAS_MEMORY, key, prompt_sig, x.shape, source=source) bad = _rescue_reference(VELOCITY_BIAS_BAD_MEMORY, key, prompt_sig, x.shape, source=source) if good is None and bad is None: if not _RESCUE_LOG["warned_no_prompt_match"]: print("[FunPack rescue] no rated trajectory for this prompt yet — rate a few gens " "(good builds the target, awful builds what to avoid). Skipping until then.") _RESCUE_LOG["warned_no_prompt_match"] = True return denoised try: thr = float(threshold) d = k_diffusion_sampling.to_d(x, sigma, denoised).detach().float() d_norm = d.norm().clamp_min(1e-8) d_unit = d / d_norm cos_good = cos_bad = None c = d_unit.clone() fired_reason = [] w = max(0.0, min(0.5, float(strength))) if good is not None: g = good.to(device=x.device, dtype=torch.float32) g_unit = g / g.norm().clamp_min(1e-8) cos_good = float(d_unit.flatten() @ g_unit.flatten()) if (1.0 - cos_good) > thr: c = c + w * (g_unit - d_unit) # pull toward good fired_reason.append(f"div_good={1.0 - cos_good:.3f}") if bad is not None: b = bad.to(device=x.device, dtype=torch.float32) b_unit = b / b.norm().clamp_min(1e-8) cu = c / c.norm().clamp_min(1e-8) cos_bad = float(cu.flatten() @ b_unit.flatten()) if cos_bad > thr: c = c - w * cos_bad * b_unit # remove bad-aligned component fired_reason.append(f"sim_bad={cos_bad:.3f}") if not fired_reason: return denoised # on-manifold w.r.t. good and clear of bad c_norm = c.norm().clamp_min(1e-8) d_corr = d_norm * (c / c_norm) denoised_corr = (x.float() - sig * d_corr).to(dtype=denoised.dtype, device=denoised.device) _RESCUE_LOG["fired"] += 1 if _RESCUE_LOG["fired"] <= 8: refs = f"good={'y' if good is not None else 'n'},bad={'y' if bad is not None else 'n'}" print(f"[FunPack rescue] sigma~{sig:.3f} {' '.join(fired_reason)} > thr={thr:.3f} " f"[{refs}] -> steered (w={w:.2f})") return denoised_corr except Exception as e: print(f"[FunPack rescue] skipped: {e}") return denoised def _apply_quality_sharpness(denoised, prev_denoised, sharpness): """E: temporal-average unsharp on the x0 prediction during the quality phase. The velocity-bias mean-pull (plus order-2 averaging) low-pass the latent and soften fine detail. This boosts the high-frequency component of the current denoised prediction relative to the previous step's prediction (a cheap temporal low-pass), restoring detail with no extra model eval. No-op when disabled, on the first quality step, or after a pulse reset (no prev). """ if not sharpness or sharpness <= 0.0 or prev_denoised is None: return denoised try: amount = max(0.0, min(1.0, float(sharpness))) prev = prev_denoised.to(device=denoised.device, dtype=denoised.dtype) if tuple(prev.shape) != tuple(denoised.shape): return denoised detail = denoised - 0.5 * (denoised + prev) # 0.5 * (denoised - prev): temporal high-pass return denoised + amount * detail except Exception as _e: _log.failed("FunPackSceneChain", "quality sharpness", _e, "this step keeps the unsharpened prediction") return denoised # --------------------------------------------------------------------------- # Blackwell (sm_120) xformers masked-attention fallback # --------------------------------------------------------------------------- # On GPUs newer than capability (9, 0), xformers' memory_efficient_attention has no # kernel that accepts a tensor attn_bias — every backend is rejected ("No operator # found ... too new ... attn_bias type is "). Unmasked attention # still works, which is why i2v ANCHOR scenes generate fine but GUIDE scenes crash: the # LTX guide path (comfy/ldm/lightricks/model.py _attention_with_guide_mask) passes a # tensor mask. PyTorch SDPA handles masks on every GPU, so when the active backend is # xformers on such a device we route only the MASKED calls to attention_pytorch via # ComfyUI's per-call optimized_attention_override hook (wrap_attn). Unmasked calls stay # on the fast xformers path. This is the programmatic equivalent of the user's # --use-pytorch-cross-attention, scoped to masked attention and installed ONLY when this # exact failure condition is detected, so other setups are untouched. def _funpack_mask_safe_attention_override(func, *args, **kwargs): # `func` is the raw backend ComfyUI selected (attention_xformers). wrap_attn has # already stamped _inside_attn_wrapper into kwargs before calling us, so invoking # attention_pytorch here cannot re-enter the override (no recursion). mask = kwargs.get("mask") if mask is None and len(args) >= 5: mask = args[4] # (q, k, v, heads, mask, ...) if mask is not None: import comfy.ldm.modules.attention as _am return _am.attention_pytorch(*args, **kwargs) return func(*args, **kwargs) def _funpack_install_mask_safe_attention(model): """When the active attention backend is xformers on a GPU too new for its tensor-bias kernels (capability > (9, 0), e.g. Blackwell sm_120), install a per-call override so MASKED attention (the LTX guide path) falls back to SDPA. No-op on supported GPUs, when ComfyUI already uses a mask-capable backend, or when another override is present. Idempotent; never raises.""" try: import comfy.ldm.modules.attention as _am if getattr(_am, "optimized_attention", None) is not getattr(_am, "attention_xformers", None): return # SDPA / sage / flash / split selected — masks already handled if not torch.cuda.is_available(): return major, minor = torch.cuda.get_device_capability() if major <= 9: return # xformers cutlass/flash tensor-bias kernels support cap <= (9, 0) to = model.model_options.setdefault("transformer_options", {}) existing = to.get("optimized_attention_override") if existing is _funpack_mask_safe_attention_override: return # already installed (idempotent across scenes/runs) if existing is not None: return # respect another tool's override; don't stomp it to["optimized_attention_override"] = _funpack_mask_safe_attention_override print(f"[FunPack AV] xformers active on capability {major}.{minor} GPU (too new for " "its tensor-bias kernels): routing MASKED attention to PyTorch SDPA so guide " "scenes don't crash. Unmasked attention stays on xformers.") except Exception as _e: print(f"[FunPack AV] mask-safe attention install skipped: {_e}") # --------------------------------------------------------------------------- # KV-Lock: variance-gated scheduler for BachVid K/V injection (arxiv 2603.09657) # # BachVid (ltx_enhancements) injects a blessed identity K/V at a fixed, # reward-scaled strength. KV-Lock makes that strength adaptive to the model's # own confidence: track how much the predicted clean latent x0 still jumps # between steps (relative step-change = a cheap proxy for the paper's variance- # of-x0 hallucination metric). High instability => the model is still drifting # => raise the multiplier so the identity K/V is held harder; once x0 settles, # the multiplier falls back toward the base strength. No DDIM inversion (LTXAV # inverts poorly) — we lock against the blessed bank, not an inverted source. # # No-op unless a BachVid K/V bank installed the shared scale list in # transformer_options. Never raises — scheduling must not break sampling. # --------------------------------------------------------------------------- _KVLOCK_TAU = 0.012 # instability threshold (relative x0 step-change) _KVLOCK_B = 2.0 # max multiplier (lets injection boost above base strength) _KVLOCK_W = 5 # sliding window length (steps) def _kvlock_find_scale_list(model): """Locate the shared kvlock_scale list build_enhancements stored in transformer_options. Returns the mutable list, or None.""" for getter in ( lambda: model.inner_model.model_patcher.model_options, lambda: model.inner_model.model_options, lambda: model.model_options, ): try: mo = getter() to = mo.get("transformer_options") if isinstance(mo, dict) else None if isinstance(to, dict): lst = to.get("funpack_kvlock_scale") if isinstance(lst, list) and lst: return lst except Exception: pass return None def _kvlock_schedule(model, denoised, prev_denoised, video_mask, state): """Update the shared kvlock_scale from x0 prediction instability. Sets the multiplier read by BachVid's K/V injector on the NEXT model() call. The list lookup is cached (sentinel "miss" = not yet searched). Never raises.""" try: if state.get("list", "miss") == "miss": state["list"] = _kvlock_find_scale_list(model) scale_list = state["list"] if scale_list is None or prev_denoised is None: return d = denoised.detach().float() p = prev_denoised.detach().float().to(d.device) if d.shape != p.shape: return if video_mask is not None and video_mask.shape[-1] == d.shape[-1]: m = video_mask.to(d.device) num = ((d - p) * m).norm() den = (d * m).norm().clamp(min=1e-6) else: num = (d - p).norm() den = d.norm().clamp(min=1e-6) instability = float(num / den) deltas = state.setdefault("deltas", []) deltas.append(instability) if len(deltas) > _KVLOCK_W: del deltas[:-_KVLOCK_W] mean_delta = sum(deltas) / len(deltas) scale_list[0] = max(0.0, min(_KVLOCK_B, mean_delta / _KVLOCK_TAU)) except Exception: pass def _get_latent_shapes(model): """Recover comfy's packed-latent stream shapes [video_shape, audio_shape, ...]. LTXAV packs its video + audio latents into one [B,1,N] tensor for sampling (comfy.utils.pack_latents: each stream flattened to [B,1,-1] and concatenated on the last dim, video first). comfy threads the per-stream shapes through the guider conds as a 'latent_shapes' CONDConstant. We read them back so FunPack perturbations can be confined to the video region. None when unavailable or single-stream. """ try: guider = getattr(model, "inner_model", None) # CFGGuider conds = getattr(guider, "conds", None) if not isinstance(conds, dict): return None for key in ("positive", "negative"): lst = conds.get(key) if not lst: continue for c in lst: mc = c.get("model_conds") if isinstance(c, dict) else None ls = mc.get("latent_shapes") if isinstance(mc, dict) else None val = getattr(ls, "cond", None) if val and len(val) > 1: return val except Exception: return None return None def _packed_video_mask(model, x): """[1,1,N] mask over a packed AV latent: 1.0 on the video stream, 0.0 on audio. Ancestral re-noising (and other video-tuned perturbations) corrupt the audio stream while helping video. With this mask we keep all FunPack steering / ancestral noise on video and let audio ride the clean deterministic flow. None when not a packed AV latent (single-stream LTXV, or layout we can't verify) -> callers then no-op. """ try: shapes = _get_latent_shapes(model) if not shapes or len(shapes) <= 1: return None if not hasattr(x, "shape") or x.dim() < 1: return None n = int(x.shape[-1]) sizes = [int(math.prod(tuple(s)[1:])) for s in shapes] if sum(sizes) != n: return None # layout doesn't match our assumption -> don't risk masking # Video = the highest-dimensional / largest stream (robust to stream ordering). video_idx = max(range(len(shapes)), key=lambda i: (len(tuple(shapes[i])), sizes[i])) mask = x.new_zeros((1, 1, n)) off = 0 for i, sz in enumerate(sizes): if i == video_idx: mask[..., off:off + sz] = 1.0 off += sz return mask except Exception: return None def _video_only(x_new, x_old, mask): """Confine a perturbation to the video region: audio keeps its base (x_old) value. Works for any perturbation expressed as new-vs-old (additive noise, steering, sharpness). For additive noise call as _video_only(x + noise, x, mask) = x + noise*mask. """ if mask is None: return x_new try: return x_old + (x_new - x_old) * mask except Exception as _e: # The mask is what keeps a video-only perturbation off the audio stream. Losing it # does not fail the render, it lets the perturbation reach the waveform. _log.failed("FunPackSceneChain", "audio-protection mask", _e, "the perturbation reached the AUDIO stream as well as the video") return x_new def _audio_clock_setup(model, x, sigmas, enabled): """Build the per-step audio-clock correction for a MiniMax H3 run, or None. Returns (factors, mask) where `factors[i]` scales the audio stream's displacement on step i and `mask` is the packed video mask. None whenever the correction cannot be made — not H3, no readable packed layout, or shifts that make it a no-op — so callers just fall through to comfy's own tangent approximation. """ if not enabled: return None try: try: from .minimax_h3 import audio_clock_factors, is_h3_model, resolve_sigma_shifts except ImportError: from minimax_h3 import audio_clock_factors, is_h3_model, resolve_sigma_shifts patcher = getattr(getattr(model, "inner_model", None), "model_patcher", None) if not is_h3_model(patcher): print("[FunPack AV] h3_audio_clock is on but this is not a MiniMax H3 model — " "the two-schedule correction it applies exists only on H3. Not running.") return None mask = _packed_video_mask(model, x) if mask is None: print("[FunPack AV] h3_audio_clock is on but the packed video+audio latent layout " "could not be read, so the audio stream cannot be located. Not running.") return None shift_v, shift_a = resolve_sigma_shifts(getattr(patcher, "model_options", None)) factors = audio_clock_factors(sigmas, shift_v, shift_a) if not factors or all(abs(f - 1.0) < 1e-6 for f in factors): print(f"[FunPack AV] h3_audio_clock is on but shift_video ({shift_v:g}) and " f"shift_audio ({shift_a:g}) put both streams on the same schedule — " f"nothing to correct. Not running.") return None worst = min(factors) print(f"[FunPack AV] h3_audio_clock on (shift_video={shift_v:g}, shift_audio={shift_a:g}) " f"— audio integrated on its own schedule over {len(factors)} step(s); the most " f"corrected step moves audio to {worst * 100.0:.0f}% of what the video grid alone " f"would have moved it.") return factors, mask except Exception as error: print(f"[FunPack AV] h3_audio_clock could not be set up ({error}) — not running.") return None def _audio_clock_step(x_new, x_old, clock, i): """Re-scale the audio stream's displacement over one step onto its own schedule. Video keeps `x_new` untouched. Audio keeps the same direction but the length the audio schedule actually calls for. Confining it to the displacement means any ancestral noise or steering the caller already restricted to video rides through unchanged. No-op without a clock, or past the end of the factor list. """ if clock is None: return x_new try: factors, mask = clock if i >= len(factors): return x_new factor = factors[i] # video region -> multiplier 1, audio region -> multiplier `factor` scale = mask + (1.0 - mask) * factor return x_old + (x_new - x_old) * scale except Exception as _e: _log.failed("FunPackSceneChain", "audio clock correction", _e, "audio integrates on the video schedule for this step (h3_audio_clock inert)") return x_new # ── the audio clock on a sampler we do not own ─────────────────────────────── # FunPack's own samplers apply the clock inside their step loop, where the step's start # and end sigma are both in hand. A stock comfy sampler has no such hook, so the only way # in is the denoised value it asks for. That works because the correction can be moved # from the step to the prediction: scaling the audio part of (x - denoised) by the same # factor makes the sampler's own euler update land the audio exactly where its schedule # says, since that update is x*r + denoised*(1-r) with r = sigma_next/sigma. # # The catch is knowing WHICH step a given call belongs to. A sampler that evaluates the # model once per step makes call index == step index, and that covers euler, res_multistep, # er_sde, dpmpp_2m and the rest of the multistep family — including the samplers ComfyUI's # own H3 template uses. A sampler that evaluates more than once per step (heun, dpm_2, # dpmpp_2s/sde) does not, and cannot be disambiguated from out here: its corrector call # lands on the same (sigma, step count) signature as the next step's predictor. # # So the wrapper checks its assumption on every call — the sigma it is handed must be the # one this step's single evaluation would use — and switches itself off, loudly, the moment # that stops holding. It never guesses. _AUDIO_CLOCK_SIGMA_TOL = 1e-4 # Measured against comfy's real samplers with a PERFECT rectified-flow predictor, so every # number is pure schedule error (scratch harness, 2026-08-06; shift 12/3, audio error as a # % of the stream's full noise->signal span): # # sampler 4 steps 8 steps 20 steps # euler 85.4 -> 0.0 38.4 -> 0.0 13.5 -> 0.0 exact, at every step count # res_multistep 68.7 -> 20.7 17.5 -> 23.0 1.0 -> 14.8 already self-corrects; clock hurts # dpmpp_2m 69.0 -> 20.2 18.2 -> 22.3 0.4 -> 14.2 same # heun/dpm_2 unchanged — two evals per step, the proxy stands down # # The lesson: this is a FIRST-ORDER integration error, so a higher-order sampler already # absorbs most of it once the schedule is fine enough — which is why res_multistep and # er_sde behaved well on H3 at 20 steps and plain-euler paths did not. At 4 steps NOTHING # absorbs it (everything sits at 68-85%) and only the clock fixes it, which is exactly the # regime turbo LoRAs put you in. # # We do not gate on this — the user picks the sampler and these are the facts they need. _AUDIO_CLOCK_EXACT_ON = ("sample_euler",) _AUDIO_CLOCK_HARMED = ("sample_res_multistep", "sample_dpmpp_2m", "sample_gradient_estimation", "sample_ipndm", "sample_ipndm_v", "sample_lms", "sample_deis") class _AudioClockDenoiser: """Denoiser proxy that applies the audio clock to what a foreign sampler asks for. Transparent for everything else: samplers reach through to `inner_model`, `latent_image` and friends, so attribute access forwards to the wrapped object. """ def __init__(self, inner, clock, sigmas): self._inner = inner self._clock = clock self._sigmas = sigmas self._step = 0 self._live = True def __getattr__(self, name): # only reached for names not on the proxy itself return getattr(self._inner, name) def __setattr__(self, name, value): if name in ("_inner", "_clock", "_sigmas", "_step", "_live"): object.__setattr__(self, name, value) else: # e.g. samplers assigning model.latent_image setattr(self._inner, name, value) def _step_for(self, sigma): """The step this call belongs to, or None if the one-eval-per-step assumption broke.""" i = self._step if i >= len(self._sigmas) - 1: return None try: want = float(self._sigmas[i]) got = float(sigma.flatten()[0]) if hasattr(sigma, "flatten") else float(sigma) except Exception: return None if abs(got - want) > _AUDIO_CLOCK_SIGMA_TOL * max(1.0, abs(want)): return None return i def __call__(self, x, sigma, **kwargs): denoised = self._inner(x, sigma, **kwargs) if not self._live: return denoised i = self._step_for(sigma) if i is None: self._live = False print("[FunPack AV] h3_audio_clock: this sampler does not evaluate the model once " "per step at the scheduled sigma, so a model call cannot be tied to a step " "from outside its loop. Correction switched OFF for the rest of this run — " "the audio is being integrated the way it was before. Use FunPack Distilled " "Flow / Hybrid Euler 2S (the clock runs inside their loop and handles this), " "or a one-eval-per-step sampler such as euler, res_multistep or er_sde.") return denoised self._step += 1 # (x - denoised) is the step's direction; scaling its audio part is the same # correction the in-loop version applies to the displacement. return _audio_clock_step(denoised, x, self._clock, i) def _audio_clock_wrap_sampler(sampler): """A SAMPLER equivalent to `sampler` with the audio clock applied around its model. Returns None when the sampler exposes no `sampler_function` to wrap — some SAMPLER producers are not KSAMPLER-shaped, and guessing at their internals is not worth it. """ fn = getattr(sampler, "sampler_function", None) if fn is None: return None def _wrapped(model, x, sigmas, extra_args=None, callback=None, disable=None, **options): clock = _audio_clock_setup(model, x, sigmas, True) if clock is None: # not H3, unreadable layout, or coincident schedules return fn(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) proxy = _AudioClockDenoiser(model, clock, sigmas) return fn(proxy, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) _wrapped.__name__ = getattr(fn, "__name__", "sampler") + "_h3_audio_clock" return comfy.samplers.KSAMPLER( _wrapped, extra_options=getattr(sampler, "extra_options", None) or {}, inpaint_options=getattr(sampler, "inpaint_options", None) or {}, ) def _video_span(model, x): """(offset, size, stream_shape) of the video stream inside a packed AV latent [B,1,N], or None when the layout can't be verified (same guards and stream heuristic as _packed_video_mask — mismatch means don't risk slicing). stream_shape is the per-stream latent shape comfy packed from (e.g. (B, C, T, H, W)), so callers can recover frame geometry from the flat span.""" try: shapes = _get_latent_shapes(model) if not shapes or len(shapes) <= 1: return None if not hasattr(x, "shape") or x.dim() < 1: return None n = int(x.shape[-1]) sizes = [int(math.prod(tuple(s)[1:])) for s in shapes] if sum(sizes) != n: return None video_idx = max(range(len(shapes)), key=lambda i: (len(tuple(shapes[i])), sizes[i])) return sum(sizes[:video_idx]), sizes[video_idx], tuple(shapes[video_idx]) except Exception: return None # Per-scene model_function_wrappers (embed guidance / score slider / output guidance / # temporal styles) live on the SHARED ModelPatcher. If a run dies between install and # restore (interrupt, OOM), the wrapper survives in-process and silently steers every # later run, stacking one deeper each time — the exact progressive-corruption failure # mode strip_funpack_block_hooks exists for, so we reuse its tag+strip pattern. _FUNPACK_SCENE_WRAPPER_TAG = "_funpack_scene_wrapper" def _make_steer_ramp(sigmas, h3): """The late-step gate every rating-driven wrapper shares, as a function of sigma. The historical gate is ``max(0, 1 - 2*sigma)``: full authority at sigma 0, nothing above 0.5. That reads sigma as "how far through the schedule are we", which holds on LTX's hand-authored schedules and is FALSE on H3. H3's schedules are ``sigma = shift*t / (1 + (shift-1)*t)`` over uniform t, so a large shift keeps sigma high right up to the final leap. Measured coverage of the old gate: shift 6 / 4 steps (turbo) 0 of 4 steps shift 12 / 12 (H3 default) 0 of 12 shift 3 / 20 4 of 20, the first two at gate 0.14 and 0.31 So embed guidance, score_slider, DynaShift and output_guidance were inert or nearly inert on H3 -- silently, since each reported itself as active. The fix reads the position on the underlying uniform base grid instead, recovered from the schedule itself rather than from a shift constant (the shift is only reliable when MiniMaxH3SigmaShift is wired, and video sampling does not require that node). For a shift-generated schedule the base grid IS the step position, so the gate becomes ``max(0, 2k/n - 1)`` -- the same "last half, ramping to full" intent the constant was written to express, at any shift and any step count. LTX is untouched: it keeps the absolute-sigma gate it was validated with. """ def legacy(sigma): return max(0.0, 1.0 - float(sigma) * 2.0) if not h3 or sigmas is None: return legacy try: vals = [float(v) for v in sigmas.flatten().tolist()] except Exception: return legacy n = len(vals) - 1 if n < 1: return legacy def progress(sigma): sigma = float(sigma) k = min(range(len(vals)), key=lambda i: abs(vals[i] - sigma)) return max(0.0, 2.0 * (k / n) - 1.0) return progress def _steer_ramp_coverage(ramp_fn, sigmas): """(gated steps, total steps, peak gate) for a schedule -- so a run can say out loud how much of it the steering actually reaches instead of only that it is 'active'.""" try: vals = [float(v) for v in sigmas.flatten().tolist()][:-1] # sigma[-1] is terminal except Exception: return None if not vals: return None gates = [ramp_fn(v) for v in vals] return sum(1 for g in gates if g > 0.0), len(gates), max(gates) if gates else 0.0 def _tag_scene_wrapper(wrapper, prev): """Mark a per-scene wrapper (and remember what it wrapped) so a later run can identify and unwind leaked ones.""" setattr(wrapper, _FUNPACK_SCENE_WRAPPER_TAG, True) setattr(wrapper, "_funpack_prev_wrapper", prev) return wrapper def _tag_funpack_hook(fn): """Mark a forward hook as FunPack's, so the run-start sweep can remove one that leaked. Anything installed on the shared diffusion model needs this: the sweep only removes what it can prove is ours, so an untagged hook is one nothing will ever clean up. """ try: from .ltx_enhancements import _FUNPACK_HOOK_TAG except ImportError: from ltx_enhancements import _FUNPACK_HOOK_TAG setattr(fn, _FUNPACK_HOOK_TAG, True) return fn def _strip_funpack_scene_wrappers(model): """Unwind FunPack per-scene wrappers leaked by a previous interrupted/failed run. Walks the recorded prev-wrapper chain back to the first non-FunPack wrapper (or none). Idempotent; never raises.""" try: w = model.model_options.get("model_function_wrapper") stripped = 0 while w is not None and getattr(w, _FUNPACK_SCENE_WRAPPER_TAG, False): w = getattr(w, "_funpack_prev_wrapper", None) stripped += 1 if stripped: if w is not None: model.model_options["model_function_wrapper"] = w else: model.model_options.pop("model_function_wrapper", None) print(f"[FunPackSceneChain] Stripped {stripped} leaked scene wrapper(s) from a previous run") return stripped except Exception: return 0 def _alg_blur_frames(model, latent_image, kappa, frame_indices=(), tail_count=0): """ALG (arXiv:2506.08456): low-pass filter selected frames of the packed video stream. I2V models over-expose the sharp anchor frame from step 0, which lets the model take a shortcut to a near-static video that just matches the reference. Bilinear downsample-then- upsample of the anchor latent (selected frames only, video stream only) at the configured factor removes the high-frequency content the shortcut needs; ALG's own per-step schedule then swaps this blurred copy back to the sharp original once sigma drops past the threshold. `frame_indices` is any iterable of frame indices within the video stream's T dimension to blur (each filtered independently) — frame 0 for the genuine i2v anchor, and/or the trailing indices of newly-appended guide-attention frames (mid_scene_guide / carry_i2v_guides-as-guide / configured guides / JoyAI memory — see EXPERIMENTAL alg_blur_guides) when extending the same idea to those. Indices outside the actual frame count are silently skipped. Returns a new tensor with the same shape as latent_image, or None on any failure/mismatch (caller then leaves ALG off — no anchor, no packed video stream, or unexpected layout). """ if latent_image is None: return None try: shapes = _get_latent_shapes(model) if not shapes: return None video_idx = max(range(len(shapes)), key=lambda i: (len(tuple(shapes[i])), int(math.prod(tuple(shapes[i])[1:])))) video_shape = tuple(shapes[video_idx]) if len(video_shape) != 5: return None # expects [B, C, T, H, W] b, c, t, h, w = (int(d) for d in video_shape) sizes = [int(math.prod(tuple(s)[1:])) for s in shapes] if sum(sizes) != int(latent_image.shape[-1]): return None # packed layout doesn't match our assumption -> don't risk it idxs = {int(i) for i in frame_indices if 0 <= int(i) < t} if tail_count > 0: idxs |= set(range(max(0, t - int(tail_count)), t)) idxs = sorted(idxs) if not idxs: return None off = sum(sizes[:video_idx]) sz = sizes[video_idx] video = latent_image[..., off:off + sz].reshape(b, c, t, h, w) dh, dw = max(1, round(h / kappa)), max(1, round(w / kappa)) video_blurred = video.clone() for idx in idxs: frame = video[:, :, idx] # [B, C, H, W] down = torch.nn.functional.interpolate(frame, size=(dh, dw), mode="bilinear", align_corners=False) blurred = torch.nn.functional.interpolate(down, size=(h, w), mode="bilinear", align_corners=False) video_blurred[:, :, idx] = blurred.to(video.dtype) out = latent_image.clone() out[..., off:off + sz] = video_blurred.reshape(b, 1, sz) return out except Exception as _e: _log.failed("FunPackSceneChain", "ALG anchor blur", _e, "the step runs on the sharp anchor — ALG de-staticking is inert") return None def _alg_prepare(model, extra_args, anchor_on, anchor_strength, tail_frames, tail_strength): """Precompute the ALG latent variants for this run. Returns ``(sharp, latents, anchor_on, tail_on)``. `latents` is keyed by ``(anchor blurred?, tail blurred?)`` — one entry per combination that can occur; the two frame sets are disjoint, so the both-blurred variant just composes the two. Anchor and guide-tail blur are fully independent (own strength, own sigma window), which is why the key is a pair rather than a single flag. Both returned flags are False when ALG cannot run at all: no real anchor (no denoise_mask), no latent_image, or a packed layout `_alg_blur_frames` refuses to guess at. Shared by the in-loop version (Distilled Flow) and the denoiser-proxy version that gives every other sampler the same behaviour — one place, so the two can't drift. """ tail_frames = max(0, int(tail_frames)) anchor_on = bool(anchor_on) tail_on = tail_frames > 0 sharp = getattr(model, "latent_image", None) if (anchor_on or tail_on) else None if not (anchor_on or tail_on) or sharp is None or (extra_args or {}).get("denoise_mask") is None: return sharp, {}, False, False anchor_blurred = _alg_blur_frames( model, sharp, max(1.0, float(anchor_strength)), frame_indices=(0,), ) if anchor_on else None tail_kappa = max(1.0, float(tail_strength)) tail_blurred = _alg_blur_frames( model, sharp, tail_kappa, tail_count=tail_frames, ) if tail_on else None both_blurred = _alg_blur_frames( model, anchor_blurred, tail_kappa, tail_count=tail_frames, ) if (anchor_blurred is not None and tail_blurred is not None) else None anchor_on = anchor_on and anchor_blurred is not None tail_on = tail_on and tail_blurred is not None latents = { (False, False): sharp, (True, False): anchor_blurred, (False, True): tail_blurred, # both → anchor → tail → None. Explicit None checks, NOT `anchor_blurred or # tail_blurred`: `or` calls bool() on the first operand, which is a multi-element # tensor when anchor blur is on but tail blur is off (both_blurred is then None) — # "Boolean value of Tensor is ambiguous". This value is only a defensive fallback: # when (True, True) is actually indexed, anchor_on and tail_on are both True, so # both_blurred is non-None and the chain never reaches the singles. (True, True): (both_blurred if both_blurred is not None else anchor_blurred if anchor_blurred is not None else tail_blurred), } return sharp, latents, anchor_on, tail_on class _ALGDenoiser: """Denoiser proxy that runs ALG's blurred/sharp swap for a sampler we can't get inside. ALG is a per-step swap of the pinned latent_image, and the swap is decided by ONE thing: the sigma of the step. That sigma is an argument of every model call, so the schedule does not need the sampler's loop at all — it only needed it in the original implementation because that is where the code happened to live. Driving it from here makes ALG work with a stock KSampler (any sampler_name), with Hybrid Euler 2S, and with multi-eval samplers such as heun, where each evaluation gets the anchor its own sigma calls for. Transparent for everything else: attribute access forwards to the wrapped denoiser. """ _OWN = ("_inner", "_latents", "_anchor_on", "_tail_on", "_anchor_thr", "_tail_thr") def __init__(self, inner, latents, anchor_on, tail_on, anchor_thr, tail_thr): self._inner = inner self._latents = latents self._anchor_on = anchor_on self._tail_on = tail_on self._anchor_thr = float(anchor_thr) self._tail_thr = float(tail_thr) def __getattr__(self, name): return getattr(self._inner, name) def __setattr__(self, name, value): if name in _ALGDenoiser._OWN: object.__setattr__(self, name, value) else: # e.g. samplers assigning model.latent_image setattr(self._inner, name, value) def __call__(self, x, sigma, **kwargs): try: s = float(sigma.flatten()[0]) if hasattr(sigma, "flatten") else float(sigma) self._inner.latent_image = self._latents[( self._anchor_on and s > self._anchor_thr, self._tail_on and s > self._tail_thr, )] except Exception: pass # never cost the step over the guidance return self._inner(x, sigma, **kwargs) def _alg_wrap_sampler(sampler, anchor_on, anchor_strength, anchor_threshold, tail_frames, tail_strength, tail_threshold): """A SAMPLER equivalent to `sampler` with ALG applied around its model. Returns None when the sampler exposes no `sampler_function` to wrap — some SAMPLER producers are not KSAMPLER-shaped, and guessing at their internals is not worth it. """ fn = getattr(sampler, "sampler_function", None) if fn is None: return None def _wrapped(model, x, sigmas, extra_args=None, callback=None, disable=None, **options): sharp, latents, a_on, t_on = _alg_prepare( model, extra_args, anchor_on, anchor_strength, tail_frames, tail_strength) if not (a_on or t_on): # no anchor, or a layout we won't guess at — run untouched return fn(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) print(f"[FunPack AV] ALG on (anchor: " f"{f'strength={anchor_strength}, sigma_threshold={anchor_threshold}' if a_on else 'off'}; " f"guide tail: " f"{f'{tail_frames} frame(s), strength={tail_strength}, sigma_threshold={tail_threshold}' if t_on else 'off'}) " f"— driven from outside {getattr(fn, '__name__', 'the sampler')}'s loop, off the " f"sigma of each model call") proxy = _ALGDenoiser(model, latents, a_on, t_on, anchor_threshold, tail_threshold) try: return fn(proxy, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) finally: model.latent_image = sharp _wrapped.__name__ = getattr(fn, "__name__", "sampler") + "_alg" return comfy.samplers.KSAMPLER( _wrapped, extra_options=getattr(sampler, "extra_options", None) or {}, inpaint_options=getattr(sampler, "inpaint_options", None) or {}, ) class _SharpenDenoiser: """Denoiser proxy that runs the quality-sharpness unsharp for a sampler we can't get inside. Same mechanism as the in-loop version on Hybrid Euler 2S / Distilled Flow: boost the x0 prediction's high-frequency component against the previous prediction. Nothing about it needs the sampler's loop — it reads only the current denoised, the previous one, and the step's sigma, and sigma is an argument of every model call. So a stock KSampler (euler, res_multistep, dpmpp_2m, …) can have the detail recovery that until now only FunPack's own samplers had. Two honest differences from the in-loop version: * On a multi-eval sampler (heun, dpmpp_2m_sde, …) "previous prediction" means the previous *evaluation*, which may be the second half of the same step rather than the step before. The high-pass is still a high-pass; its magnitude is not identical. * `prev` holds the SHARPENED result, matching the in-loop samplers, where `prev_denoised = denoised` runs after the sharpen. That makes the boost mildly self-limiting instead of compounding. Transparent for everything else: attribute access forwards to the wrapped denoiser, so this composes with the ALG proxy in either order. """ _OWN = ("_inner", "_amount", "_thr", "_mask", "_prev", "_mask_done") def __init__(self, inner, amount, threshold, mask): self._inner = inner self._amount = float(amount) self._thr = float(threshold) self._mask = mask self._prev = None def __getattr__(self, name): return getattr(self._inner, name) def __setattr__(self, name, value): if name in _SharpenDenoiser._OWN: object.__setattr__(self, name, value) else: # e.g. samplers (or the ALG proxy) assigning latent_image setattr(self._inner, name, value) def __call__(self, x, sigma, **kwargs): denoised = self._inner(x, sigma, **kwargs) try: s = float(sigma.flatten()[0]) if hasattr(sigma, "flatten") else float(sigma) except Exception as _e: # noqa: BLE001 # Once, not per evaluation — but never nothing. Silently returning here would # leave quality_sharpness switched on, reported as active, and doing nothing for # the whole run. _log.failed("FunPackStudio", "quality sharpness (unreadable sigma)", _e, "this sampler's steps are NOT being sharpened") return denoised if s > self._thr: # Outside the window: nothing to sharpen, but the prediction is still KEPT. The # in-loop samplers carry prev_denoised across the phase boundary, so the first # quality step sharpens against the last pre-quality one; dropping it here would # make the same setting mean something different on a KSampler. The cost is one # latent (~40 MB on a 97-frame 768x512 H3 clip) for the length of the run, which # is not a memory problem worth breaking parity over. try: self._prev = denoised.detach() except Exception: # noqa: BLE001 self._prev = None return denoised try: denoised = _video_only( _apply_quality_sharpness(denoised, self._prev, self._amount), denoised, self._mask) except Exception as _e: # noqa: BLE001 _log.failed("FunPackStudio", "quality sharpness", _e, "this evaluation keeps the unsharpened prediction") try: self._prev = denoised.detach() except Exception: # noqa: BLE001 self._prev = None return denoised def release(self): """Drop the retained prediction. Called when the sampler returns, so the tensor is gone before the decode asks for its peak rather than waiting on garbage collection.""" self._prev = None self._mask = None def _sharpen_wrap_sampler(sampler, sharpness, start_pct): """A SAMPLER equivalent to `sampler` with quality sharpness applied around its model. Returns `sampler` unchanged when sharpness is off, and None when the sampler exposes no `sampler_function` to wrap. """ sharpness = max(0.0, min(1.0, float(sharpness or 0.0))) if sharpness <= 0.0: return sampler fn = getattr(sampler, "sampler_function", None) if fn is None: return None def _wrapped(model, x, sigmas, extra_args=None, callback=None, disable=None, **options): # The window is a fraction of the schedule, converted to the sigma that starts it — # exactly how Hybrid Euler 2S derives its quality phase from high_quality_pct, so the # same number means the same thing on both samplers. try: sched_steps = max(1, int(sigmas.shape[0]) - 1) late_start = _get_late_start_index(sched_steps, start_pct) thr = float(sigmas[late_start].item()) if late_start < int(sigmas.shape[0]) else None except Exception: thr = None if thr is None: return fn(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) mask = _packed_video_mask(model, x) print(f"[FunPack] quality sharpness {sharpness:.2f} on the last " f"{start_pct * 100:.0f}% of the schedule (sigma <= {thr:.4f}) — driven from " f"outside {getattr(fn, '__name__', 'the sampler')}'s loop, off the sigma of " f"each model call" f"{'' if mask is not None else '; single-stream latent, no audio to protect'}") proxy = _SharpenDenoiser(model, sharpness, thr, mask) try: return fn(proxy, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) finally: proxy.release() _wrapped.__name__ = getattr(fn, "__name__", "sampler") + "_sharpen" return comfy.samplers.KSAMPLER( _wrapped, extra_options=getattr(sampler, "extra_options", None) or {}, inpaint_options=getattr(sampler, "inpaint_options", None) or {}, ) class _JoyAIMemoryBank: """JoyAI-Echo cross-shot memory: a rolling set of clean prior-shot latent frames. Mirrors PairedAudioVideoMemoryBank's policy (JoyAI_Echo configs/inference.yaml): the first ``num_fix`` entries are pinned permanently (the opening shots = global story anchor); the rest is a most-recent window so the total never exceeds ``max_size``. Each entry is a paired (video_frame, audio_frame) of clean latents from a finished scene; the audio half is None until joyai_audio_memory is on. Pure bookkeeping — injection is done by the sampler (video via LTX guide attention, audio via a protected prefix), so this holds no model state and is unit-testable. """ def __init__(self, max_size=7, num_fix=3): self.max_size = max(1, int(max_size)) self.num_fix = max(0, min(int(num_fix), self.max_size)) self.entries = [] # list of (video_frame, audio_frame) tuples, oldest first def add(self, video_frame, audio_frame=None): if video_frame is None: return self.entries.append((video_frame, audio_frame)) if len(self.entries) <= self.max_size: return fixed = self.entries[:self.num_fix] recent = self.entries[self.num_fix:] room = self.max_size - len(fixed) self.entries = fixed + (recent[-room:] if room > 0 else []) def frames(self): """Video latent frames, oldest first (the injection order for guide attention).""" return [v for v, _ in self.entries] def audio(self): """Paired audio latent frames aligned 1:1 with frames(); entries hold None when audio is off.""" return [a for _, a in self.entries] def _find_schedule_anchor_index(sigmas, total_steps, schedule_progress): if sigmas is None or total_steps <= 1: return 0 schedule_progress = max(0.0, min(1.0, schedule_progress)) return min(total_steps - 1, max(0, int(round(schedule_progress * max(0, total_steps - 1))))) def _resolve_motion_pulse_options(motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength): mode = (motion_pulse_mode or "off").lower() if mode not in MOTION_PULSE_MODES: mode = "off" if mode == "off": return { "enabled": False, "start_pct": 0.30, "count": 0, "spacing_pct": 0.22, "strength": 0.0, "noise": 0.0, } start_pct = 0.30 if motion_pulse_start_pct is None else float(motion_pulse_start_pct) spacing_pct = 0.22 if motion_pulse_spacing_pct is None else float(motion_pulse_spacing_pct) strength = 0.85 if motion_pulse_strength is None else float(motion_pulse_strength) count = 2 if motion_pulse_count is None else int(motion_pulse_count) start_pct = max(0.02, min(0.90, start_pct)) spacing_pct = max(0.04, min(0.45, spacing_pct)) strength = max(0.0, min(1.0, strength)) count = max(1, min(6, count)) if mode == "balanced": count = min(count, 1) strength = 0.55 if motion_pulse_strength is None else min(strength, 0.70) elif mode == "aggressive": count = max(2, count) strength = max(strength, 0.85) noise = 0.10 + strength * 0.55 return { "enabled": True, "start_pct": start_pct, "count": count, "spacing_pct": spacing_pct, "strength": strength, "noise": max(0.0, min(0.80, noise)), } def _get_late_start_index(total_steps, high_quality_pct): high_quality_pct = max(0.0, min(1.0, float(high_quality_pct))) late_steps = max(1, int(math.ceil(total_steps * high_quality_pct))) if high_quality_pct > 0.0 else 0 return max(0, total_steps - late_steps) def _build_motion_pulse_steps(sigmas, total_steps, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength): options = _resolve_motion_pulse_options( motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, ) if not options["enabled"] or total_steps <= 2: return [], options late_start = _get_late_start_index(total_steps, high_quality_pct) latest_anchor_pct = max(0.04, min(0.92, (late_start - 1) / max(1, total_steps - 1))) pulse_steps = [] used_anchors = set() for pulse_index in range(options["count"]): trigger_pct = options["start_pct"] + options["spacing_pct"] * pulse_index if trigger_pct >= latest_anchor_pct: break anchor_index = _find_schedule_anchor_index(sigmas, total_steps, trigger_pct) anchor_index = min(total_steps - 1, max(0, anchor_index)) if anchor_index in used_anchors or anchor_index >= late_start: continue used_anchors.add(anchor_index) pulse_steps.append({ "step_index": anchor_index, "noise": options["noise"], }) return pulse_steps, options def _prepare_dynamic_sigmas(sigmas, high_quality_pct, motion_pulse_mode="off", motion_pulse_start_pct=0.30, motion_pulse_count=2, motion_pulse_spacing_pct=0.22, motion_pulse_strength=0.85): if sigmas is None or not isinstance(sigmas, torch.Tensor): return None, None, [], 0.0 base_sigmas = sigmas.detach().clone() total_steps = max(0, int(base_sigmas.shape[0]) - 1) if total_steps <= 0: return base_sigmas, None, [], 0.0 late_start = _get_late_start_index(total_steps, high_quality_pct) quality_sigma_start = None if late_start < base_sigmas.shape[0]: quality_sigma_start = float(base_sigmas[late_start].item()) pulse_steps, motion_pulse_options = _build_motion_pulse_steps( base_sigmas, total_steps, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, ) return base_sigmas, quality_sigma_start, pulse_steps, motion_pulse_options["noise"] def _order2_ancestral_denoised(denoised, prev_denoised, h, prev_h): """ Linear extrapolation of the denoised estimate using the previous step's value. Equivalent to the DPM-Solver++ 2M approach applied to the ancestral phase. Gives second-order accuracy at zero extra model-call cost. """ if prev_denoised is None or prev_h is None or prev_h < 1e-7 or h < 1e-7: return denoised r = max(0.25, min(4.0, prev_h / h)) c1 = 1.0 + 0.5 / r c2 = 0.5 / r try: extrap = c1 * denoised - c2 * prev_denoised.to(device=denoised.device, dtype=denoised.dtype) return extrap except Exception: return denoised def _sample_const_rf_full(model, x, sigmas, extra_args, callback, disable, eta, s_noise, eta_final, high_quality_pct, correction_blend, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, motion_pulse_noise, motion_pulse_steps, velocity_bias_mode, velocity_bias_strength, 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, h3_audio_clock=False): """Full-feature rectified-flow sampler for CONST models (LTXAV). Rectified-flow-correct port of the hybrid sampler so its features actually run on LTXAV (which stock comfy routes to euler_ancestral_RF, bypassing the custom loop): - early phase: ancestral RF euler with eta decay, AB2 (order-2) denoised extrapolation, and anti-stiffness motion pulses - late (quality) phase: deterministic Heun corrector — the RF-correct 2nd order; DPM++ 2S is eps-only and intentionally not used here — with the hybrid's progressive correction_blend schedule - velocity-bias capture/apply and reactive rescue around every model eval With motion/velocity/rescue off AND high_quality_pct=0 it reduces to stock euler_ancestral_RF. The base RF step matches comfy.sample_euler_ancestral_RF. """ extra_args = {} if extra_args is None else extra_args seed = extra_args.get("seed", None) noise_sampler = k_diffusion_sampling.default_noise_sampler(x, seed=seed) try: s_noise = s_noise * getattr( model.inner_model.model_patcher.get_model_object('model_sampling'), "noise_scale", 1.0) except Exception: pass total_steps = max(0, len(sigmas) - 1) if total_steps <= 0: return x high_quality_pct = max(0.0, min(1.0, float(high_quality_pct))) correction_blend = max(0.0, min(1.0, float(correction_blend))) eta_final = max(0.0, min(float(eta), float(eta_final))) sched_steps = max(1, int(len(sigmas)) - 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, 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 = { int(item.get("step_index", -1)): max(0.0, float(item.get("noise", motion_pulse_noise))) for item in (motion_pulse_steps or []) if isinstance(item, dict) } late_start = _get_late_start_index(sched_steps, high_quality_pct) quality_sigma_start = float(sigmas[late_start].item()) if late_start < sigmas.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 _kvlock_state = {} # Audio-safe sampling: on a packed LTXAV latent, keep ancestral noise + steering on the # video stream and let audio ride the clean deterministic flow (ancestral re-noising # corrupts audio). None for single-stream LTXV -> all of this is a no-op. video_mask = _packed_video_mask(model, x) if video_mask is not None: n_aud = int((video_mask < 0.5).sum().item()) print(f"[FunPack AV] packed audio+video latent detected -> audio-safe sampling " f"(audio held deterministic on {n_aud} of {video_mask.shape[-1]} packed dims)") normalize_strength = max(0.0, min(1.0, float(normalize_strength))) ref_scale = [None] # video-only latent-normalization reference (anti-overbake), opt-in if normalize_strength > 0.0: print(f"[FunPack AV] latent normalization on (strength={normalize_strength}, " f"start_sigma={normalize_start_sigma}, video-only)") audio_clock = _audio_clock_setup(model, x, sigmas, h3_audio_clock) for i in comfy.utils.model_trange(total_steps, disable=disable): sigma = sigmas[i] sigma_next = sigmas[i + 1] in_quality = quality_sigma_start is not None and float(sigma.item()) <= quality_sigma_start velocity_target = _velocity_bias_target(sigmas, sigma) if not in_quality: pulse = motion_step_noise.get(int(i), 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 prev_h = None 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), prompt_sig=rescue_prompt_sig, source=velocity_bias_source) x = _video_only(x, x_pre, video_mask) denoised = model(x, sigma * s_in, **extra_args) _kvlock_schedule(model, denoised, prev_denoised, video_mask, _kvlock_state) if rescue_mode or _velocity_bias_enabled(velocity_bias_mode, "capture"): _capture_velocity_bias(velocity_refinement_key, velocity_target, x, sigma, denoised, prompt_sig=rescue_prompt_sig) if rescue_mode and velocity_target is not None: denoised = _video_only(_rescue_denoised( denoised, x, sigma, velocity_refinement_key, velocity_target, rescue_threshold, rescue_strength, prompt_sig=rescue_prompt_sig, source=velocity_bias_source, ), denoised, video_mask) # E: restore high-frequency detail lost to the velocity-bias mean-pull (quality phase only). if in_quality: 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 RF loop. denoised = _normalize_video_denoised( denoised, video_mask, sigma, ref_scale, normalize_strength, normalize_start_sigma, ) if callback is not None: callback({'x': x, 'i': i, 'sigma': sigma, 'sigma_hat': sigma, 'denoised': denoised}) h = float((sigma - sigma_next).abs().item()) # AB2 (order-2) extrapolation on the x0/denoised sequence — RF-valid. if prev_denoised is not None and prev_h is not None and prev_h > 1e-7 and h > 1e-7: r = max(0.1, min(5.0, h / prev_h)) try: denoised_eff = (1.0 + r / 2.0) * denoised - (r / 2.0) * prev_denoised.to(device=denoised.device, dtype=denoised.dtype) except Exception: denoised_eff = denoised else: denoised_eff = denoised # Audio rides plain euler — exclude it from AB2 (2nd-order extrapolation corrupts audio). denoised_eff = _video_only(denoised_eff, denoised, video_mask) prev_denoised = denoised.detach() prev_h = h if sigma_next == 0: # Identical displacement to the euler-RF step below at sigma_next = 0, so the # audio clock corrects this final step like any other. x = _audio_clock_step(denoised_eff, x, audio_clock, i) continue if in_quality: # Deterministic Heun corrector (RF-correct 2nd order). Progressive blend: # first half of quality steps = euler, second half = full Heun, matching # the hybrid sampler's correction schedule. Note: to_d-based euler is # identical to the RF flow update, so this is consistent with the early phase. if num_quality_steps <= 1: effective_blend = correction_blend else: q_idx = max(0, i - 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: # Audio in the predicted state gets the clock too: the corrector's audio # output is discarded, but video and audio share one packed sequence, so # an over-stepped audio half skews the video correction via joint attention. x_pred = _audio_clock_step(x + d1 * dt, x, audio_clock, i) denoised_pred = model(x_pred, sigma_next * s_in, **extra_args) d2 = k_diffusion_sampling.to_d(x_pred, sigma_next, denoised_pred) d_use = d1 + effective_blend * ((d1 + d2) * 0.5 - d1) else: d_use = d1 # Audio rides the plain euler direction (d1); only video gets the Heun correction. d_use = _video_only(d_use, d1, video_mask) x = _audio_clock_step(x + d_use * dt, x, audio_clock, i) # Heun changed x with a corrected direction; invalidate AB2 history. prev_denoised = None prev_h = None 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) # 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. er = sigma_next / sigma x_det = er * x + (1 - er) * denoised_eff if effective_eta > 0: downstep_ratio = 1 + (sigma_next / sigma - 1) * effective_eta sigma_down = sigma_next * downstep_ratio alpha_ip1 = 1 - sigma_next alpha_down = 1 - sigma_down sigma_down_i_ratio = sigma_down / sigma x_anc = sigma_down_i_ratio * x + (1 - sigma_down_i_ratio) * denoised_eff renoise_coeff = (sigma_next ** 2 - sigma_down ** 2 * alpha_ip1 ** 2 / alpha_down ** 2) ** 0.5 x_anc = (alpha_ip1 / alpha_down) * x_anc + noise_sampler(sigma, sigma_next) * s_noise * renoise_coeff # Video gets full ancestral noise; audio stays on the deterministic step. # The clock is applied to the combined result: its audio half is exactly # x_det's (the noise above is video-only), so it scales the deterministic # audio displacement and leaves the noised video untouched. x = _audio_clock_step(_video_only(x_anc, x_det, video_mask), x, audio_clock, i) else: x = _audio_clock_step(x_det, x, audio_clock, i) return x def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1.0, high_quality_pct=0.35, correction_blend=1.0, quality_sharpness=0.0, quality_sigma_start=None, motion_pulse_mode="off", motion_pulse_start_pct=0.30, motion_pulse_count=2, motion_pulse_spacing_pct=0.22, motion_pulse_strength=0.85, motion_pulse_noise=0.0, motion_pulse_steps=None, velocity_bias_mode="off", velocity_bias_strength=0.0, velocity_bias_source="mean", velocity_refinement_key="default", rescue_mode=False, rescue_threshold=0.15, rescue_strength=0.2, rescue_prompt_sig=None, eta_final=1.0, normalize_strength=0.0, normalize_start_sigma=0.9, h3_audio_clock=False): """ Hybrid sampler: - Early schedule: Euler ancestral with order-2 denoised extrapolation for motion/anatomy buildup. Order-2 reuses the previous step's denoised to extrapolate the score direction, giving DPM-Solver++ 2M accuracy at zero extra model-call cost. - Late schedule: deterministic DPM-Solver++(2S) ODE refinement for detail, with progressive correction_blend — first half of quality steps use single- eval Euler ODE, second half use the full configured 2S correction. This cuts quality-phase model calls by roughly half while preserving the 2S benefit where sigma is lowest and it matters most. - Eta decay: when eta_final < eta, ancestral noise strength decays toward eta_final as sigma approaches the quality boundary, giving a cleaner transition into deterministic refinement. - Motion pulses: optional monotonic noise kicks before the late quality phase. """ # Normalize velocity/rescue params up front — both the CONST (rectified-flow) # branch and the main hybrid loop rely on them. velocity_bias_mode = (velocity_bias_mode or "off").lower() if velocity_bias_mode not in VELOCITY_BIAS_MODES: velocity_bias_mode = "off" velocity_bias_strength = max(0.0, min(3.0, float(velocity_bias_strength or 0.0))) velocity_bias_source = (velocity_bias_source or "mean").lower() if velocity_bias_source not in ("mean", "nearest"): velocity_bias_source = "mean" rescue_mode = bool(rescue_mode) rescue_threshold = max(0.0, min(1.0, float(rescue_threshold or 0.0))) rescue_strength = max(0.0, min(0.5, float(rescue_strength or 0.0))) rescue_prompt_sig = _normalize_sig(rescue_prompt_sig) # capture also uses this if rescue_mode: _RESCUE_LOG["warned_no_memory"] = False _RESCUE_LOG["warned_no_prompt_match"] = False _RESCUE_LOG["fired"] = 0 if bool(rescue_mode): # One-shot startup diagnostic so it is unambiguous why rescue does or does # not fire: confirms the option reached the sampler, whether this model # bypasses the custom loop (CONST), and which steps are eligible. try: _is_const = isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST) except Exception: _is_const = "unknown" try: _ts = max(0, len(sigmas) - 1) _elig = [(i, round(float(sigmas[i].item()), 4)) for i in range(_ts) if _velocity_bias_target(sigmas, sigmas[i]) is not None] except Exception: _elig = [] print(f"[FunPack rescue] requested — strength={float(rescue_strength)}, " f"threshold={float(rescue_threshold)}, " f"prompt_sig={'yes' if isinstance(rescue_prompt_sig, torch.Tensor) else 'no'}, " f"model_is_CONST={_is_const}, eligible_steps={_elig}") if _is_const is True: print("[FunPack rescue] CONST model -> using rectified-flow rescue path (velocity/rescue active).") if not _elig: print("[FunPack rescue] WARNING: no sigma step matches a velocity target " "(normalized ~0.90/0.80 +/-0.065) -> rescue/capture cannot fire on this schedule.") if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST): # CONST (rectified-flow, e.g. LTXAV): stock comfy routes to euler_ancestral_RF, # bypassing the entire custom loop. Run the RF-correct full-feature port so the # hybrid sampler's behavior (motion pulses, order-2, quality correction, eta # decay, velocity bias, rescue) actually applies on LTXAV. return _sample_const_rf_full( model, x, sigmas, extra_args, callback, disable, eta, s_noise, eta_final, high_quality_pct, correction_blend, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, motion_pulse_noise, motion_pulse_steps, velocity_bias_mode, velocity_bias_strength, velocity_refinement_key, rescue_mode, rescue_threshold, rescue_strength, rescue_prompt_sig, quality_sharpness=quality_sharpness, velocity_bias_source=velocity_bias_source, normalize_strength=normalize_strength, normalize_start_sigma=normalize_start_sigma, h3_audio_clock=h3_audio_clock, ) # Below here is the eps-parameterised loop, which MiniMax H3 (CONST) never reaches — # h3_audio_clock is deliberately not wired into it rather than added as dead code. extra_args = {} if extra_args is None else extra_args seed = extra_args.get("seed", None) noise_sampler = k_diffusion_sampling.default_noise_sampler(x, seed=seed) total_steps = max(0, len(sigmas) - 1) if total_steps <= 0: return x high_quality_pct = max(0.0, min(1.0, float(high_quality_pct))) correction_blend = max(0.0, min(1.0, float(correction_blend))) eta_final = max(0.0, min(float(eta), float(eta_final))) if not motion_pulse_steps: _, _, motion_pulse_steps, computed_motion_pulse_noise = _prepare_dynamic_sigmas( sigmas, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, ) motion_pulse_noise = computed_motion_pulse_noise motion_pulse_noise = max(0.0, float(motion_pulse_noise)) motion_step_noise = { int(item.get("step_index", -1)): max(0.0, float(item.get("noise", motion_pulse_noise))) for item in (motion_pulse_steps or []) if isinstance(item, dict) } s_in = x.new_ones([x.shape[0]]) callback_step = 0 # Resolve quality phase boundary late_start = _get_late_start_index(total_steps, high_quality_pct) if quality_sigma_start is None: if late_start < sigmas.shape[0]: quality_sigma_start = float(sigmas[late_start].item()) else: quality_sigma_start = float(quality_sigma_start) num_quality_steps = total_steps - late_start # Order-2 ancestral state prev_denoised = None prev_h = None quality_step_index = 0 _kvlock_state = {} video_mask = _packed_video_mask(model, x) for i in comfy.utils.model_trange(total_steps, disable=disable): sigma = sigmas[i] sigma_next = sigmas[i + 1] in_quality_phase = quality_sigma_start is not None and float(sigma.item()) <= quality_sigma_start # Velocity-bias capture and rescue must run wherever a velocity target lands, # regardless of phase. With short distilled schedules the only target step # (norm sigma ~0.9) often falls inside the quality phase, so gating these to # the early branch alone would silently disable them. velocity_target = _velocity_bias_target(sigmas, sigma) if not in_quality_phase: # Adaptive eta: decay from eta toward eta_final as sigma # approaches the quality boundary. effective_eta = _effective_eta(eta, eta_final, sigmas, sigma) pulse_noise = motion_step_noise.get(int(i), 0.0) if pulse_noise > 0.0: x = _apply_motion_pulse(x, sigma, sigma_next, pulse_noise, noise_sampler) # Pulse modifies x; previous denoised is no longer a valid # second-order estimate for the next step. prev_denoised = None prev_h = None if _velocity_bias_enabled(velocity_bias_mode, "apply"): x = _apply_velocity_bias(x, velocity_refinement_key, velocity_target, velocity_bias_strength, sigma_ratio=_sigma_ratio(sigmas, sigma), prompt_sig=rescue_prompt_sig, source=velocity_bias_source) denoised = model(x, sigma * s_in, **extra_args) _kvlock_schedule(model, denoised, prev_denoised, video_mask, _kvlock_state) if rescue_mode or _velocity_bias_enabled(velocity_bias_mode, "capture"): _capture_velocity_bias(velocity_refinement_key, velocity_target, x, sigma, denoised, prompt_sig=rescue_prompt_sig) if rescue_mode and velocity_target is not None: denoised = _rescue_denoised( denoised, x, sigma, velocity_refinement_key, velocity_target, rescue_threshold, rescue_strength, prompt_sig=rescue_prompt_sig, source=velocity_bias_source, ) if callback is not None: callback({ "x": x, "i": callback_step, "sigma": sigma, "sigma_hat": sigma, "denoised": denoised, }) callback_step += 1 h = float((sigma - sigma_next).abs().item()) denoised_eff = _order2_ancestral_denoised(denoised, prev_denoised, h, prev_h) sigma_down, sigma_up = k_diffusion_sampling.get_ancestral_step(sigma, sigma_next, eta=effective_eta) if sigma_down == 0: x = denoised_eff else: d = k_diffusion_sampling.to_d(x, sigma, denoised_eff) dt = sigma_down - sigma x = x + d * dt if sigma_next > 0 and effective_eta > 0 and s_noise > 0: x = x + noise_sampler(sigma, sigma_next) * s_noise * sigma_up prev_denoised = denoised.detach() prev_h = h else: # Quality phase: progressive correction_blend. # First half of quality steps use blend=0 (single-eval Euler ODE), # second half use the configured blend (full 2S correction). # 2S matters most at the lowest sigmas, so this concentrates the # expensive second model call where it has the most impact. if num_quality_steps <= 1: effective_blend = correction_blend else: mid_quality = num_quality_steps // 2 effective_blend = 0.0 if quality_step_index < mid_quality else correction_blend if _velocity_bias_enabled(velocity_bias_mode, "apply"): x = _apply_velocity_bias(x, velocity_refinement_key, velocity_target, velocity_bias_strength, sigma_ratio=_sigma_ratio(sigmas, sigma), prompt_sig=rescue_prompt_sig, source=velocity_bias_source) denoised = model(x, sigma * s_in, **extra_args) _kvlock_schedule(model, denoised, prev_denoised, video_mask, _kvlock_state) if rescue_mode or _velocity_bias_enabled(velocity_bias_mode, "capture"): _capture_velocity_bias(velocity_refinement_key, velocity_target, x, sigma, denoised, prompt_sig=rescue_prompt_sig) if rescue_mode and velocity_target is not None: denoised = _rescue_denoised( denoised, x, sigma, velocity_refinement_key, velocity_target, rescue_threshold, rescue_strength, prompt_sig=rescue_prompt_sig, source=velocity_bias_source, ) # E: restore high-frequency detail lost to the velocity-bias mean-pull. denoised = _apply_quality_sharpness(denoised, prev_denoised, quality_sharpness) if callback is not None: callback({ "x": x, "i": callback_step, "sigma": sigma, "sigma_hat": sigma, "denoised": denoised, }) callback_step += 1 x, _ = _hybrid_ode_step( model, x, sigma, sigma_next, s_in, extra_args, effective_blend, denoised=denoised, ) quality_step_index += 1 return x class FunPackHybridEuler2SSampler: @classmethod def INPUT_TYPES(cls): return { "required": { "eta": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Ancestral stochasticity at the start of sampling. Keep at 1.0 for classic ancestral behaviour." }), "eta_final": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Eta value at the quality phase boundary. When below eta, ancestral noise decays linearly toward this value as sigma approaches the quality phase. Lower values give a cleaner hand-off into deterministic refinement. Set equal to eta to disable decay." }), "s_noise": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Noise scale for ancestral noise injection." }), "high_quality_pct": ("FLOAT", { "default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Fraction of late denoise steps that enter the quality phase. The first half of quality steps use single-eval Euler ODE; the second half use the full 2S correction." }), "correction_blend": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Blend between Euler ODE (0.0) and 2S correction (1.0) for the second half of quality-phase steps." }), "quality_sharpness": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Restores fine detail lost to velocity-bias softening. Temporal-average unsharp on the x0 prediction during the quality phase only. 0 disables; 0.2-0.4 typical when velocity bias is on. Free (no extra model eval)." }), "motion_pulse_mode": (MOTION_PULSE_MODES, { "default": "off", "tooltip": "Adds early/mid anti-stiffness motion pulses. Off preserves legacy sampler behavior." }), "motion_pulse_start_pct": ("FLOAT", { "default": 0.30, "min": 0.0, "max": 0.95, "step": 0.01, "tooltip": "Sampling progress where the first anti-stiffness pulse is applied." }), "motion_pulse_count": ("INT", { "default": 2, "min": 1, "max": 6, "step": 1, "tooltip": "How many early/mid motion pulses to request before the late quality phase." }), "motion_pulse_spacing_pct": ("FLOAT", { "default": 0.22, "min": 0.04, "max": 0.45, "step": 0.01, "tooltip": "Progress spacing between motion pulses." }), "motion_pulse_strength": ("FLOAT", { "default": 0.85, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "How strongly motion pulses add monotonic noise kicks. Higher values push harder against stale image references." }), "velocity_bias_mode": (VELOCITY_BIAS_MODES, { "default": "off", "tooltip": "Experimental: capture/apply averaged early model velocity around normalized sigma 0.9 and 0.8. Off preserves legacy sampler behavior." }), "velocity_bias_strength": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 3.0, "step": 0.05, "tooltip": "Strength of the remembered velocity (action) injected at the structure sigma. 0 disables. ~0.15 = subtle spice; 0.3-1.0 = clear action crossover; 2-3 approaches full action replacement (capped so the current gen isn't wiped). Creative tool, not for consistency." }), "velocity_bias_source": (["mean", "nearest"], { "default": "mean", "tooltip": "How velocity bias / rescue pick a good direction. 'mean' = prompt-blind global average (legacy). 'nearest' = single best-matching prompt cluster — preserves one real good gen's detail instead of a washed-out average (less softening). Affects both apply and rescue." }), "velocity_refinement_key": ("STRING", { "default": "default", "multiline": False, "tooltip": "Memory key used to capture/apply early velocity bias." }), "rescue_mode": ("BOOLEAN", { "default": False, "tooltip": "Reactive in-flight rescue, rating-gated. Steers each eligible step toward trajectories you rated good and away from ones you rated Awful (matched to the current prompt). Learns automatically from ratings while on — no separate capture step needed. A no-op until you've rated a few gens for this prompt/key (a positive rating builds the target, an Awful builds what to avoid)." }), "rescue_threshold": ("FLOAT", { "default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Fires when the step has diverged from the good trajectory by more than this (1 - cosine) OR aligned with a bad trajectory by more than this (cosine). Lower = corrects more eagerly. 0.10-0.20 typical; raise toward 0.4+ for only severe cases." }), "rescue_strength": ("FLOAT", { "default": 0.2, "min": 0.0, "max": 0.5, "step": 0.01, "tooltip": "How hard to pull toward good / push away from bad when triggered (magnitude preserved, no energy injected). Keep moderate; 0.5 is a strong correction." }), }, "optional": { "sigmas": ("SIGMAS",), # Appended last on purpose so existing saved nodes keep their widget order. "normalize_strength": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Video-only latent normalization (anti-overbake / oversaturation / colour drift) stacked on this RF loop. 0 = off. 0.5 = gentle. Audio is never touched (LTXAV/CONST path). ~zero overhead.", }), "normalize_start_sigma": ("FLOAT", { "default": 0.9, "min": 0.0, "max": 1.0, "step": 0.025, "tooltip": "Sigma at/below which latent normalization activates and anchors its reference. Only used when normalize_strength > 0.", }), } } RETURN_TYPES = ("SAMPLER", "SIGMAS") RETURN_NAMES = ("sampler", "sigmas") FUNCTION = "get_sampler" CATEGORY = "FunPack/Sampling" DESCRIPTION = ( "Hybrid sampler: early Euler ancestral with order-2 denoised extrapolation for motion, " "late DPM-Solver++(2S) ODE for quality with progressive correction blending. " "Optional eta decay, anti-stiffness motion pulses, velocity bias, and reactive rescue. " "On rectified-flow models (CONST, e.g. LTXAV) it runs an RF-correct port of the same " "features (Heun corrector in place of 2S) instead of falling back to plain euler-ancestral." ) def get_sampler(self, eta, s_noise, high_quality_pct, correction_blend, quality_sharpness=0.0, motion_pulse_mode="off", motion_pulse_start_pct=0.30, motion_pulse_count=2, motion_pulse_spacing_pct=0.22, motion_pulse_strength=0.85, velocity_bias_mode="off", velocity_bias_strength=0.0, velocity_bias_source="mean", velocity_refinement_key="default", rescue_mode=False, rescue_threshold=0.15, rescue_strength=0.2, rescue_prompt_sig=None, sigmas=None, eta_final=1.0, normalize_strength=0.0, normalize_start_sigma=0.9): prepared_sigmas, quality_sigma_start, motion_pulse_steps, motion_pulse_noise = _prepare_dynamic_sigmas( sigmas, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, ) sampler = comfy.samplers.KSAMPLER( sample_funpack_hybrid_euler_2s, extra_options={ "eta": eta, "s_noise": s_noise, "high_quality_pct": high_quality_pct, "correction_blend": correction_blend, "quality_sharpness": quality_sharpness, "quality_sigma_start": quality_sigma_start, "motion_pulse_mode": motion_pulse_mode, "motion_pulse_start_pct": motion_pulse_start_pct, "motion_pulse_count": motion_pulse_count, "motion_pulse_spacing_pct": motion_pulse_spacing_pct, "motion_pulse_strength": motion_pulse_strength, "motion_pulse_noise": motion_pulse_noise, "motion_pulse_steps": motion_pulse_steps, "velocity_bias_mode": velocity_bias_mode, "velocity_bias_strength": velocity_bias_strength, "velocity_bias_source": velocity_bias_source, "velocity_refinement_key": velocity_refinement_key, "rescue_mode": rescue_mode, "rescue_threshold": rescue_threshold, "rescue_strength": rescue_strength, "rescue_prompt_sig": rescue_prompt_sig, "eta_final": eta_final, "normalize_strength": normalize_strength, "normalize_start_sigma": normalize_start_sigma, } ) return (sampler, prepared_sigmas) def _rf_ancestral(model): """True when this model is a rectified-flow (CONST) one — H3 and LTXAV both are. Decides which ancestral formulation is correct; a wrong answer only changes how noise is added, never whether sampling runs, so an unreadable model falls back to the eps form. """ try: return isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST) except Exception: return False def sample_funpack_distilled_flow(model, x, sigmas, extra_args=None, callback=None, disable=None, order=2, s_noise=0.0, final_correction_steps=0, ab2_ramp=False, normalize_strength=0.0, normalize_start_sigma=0.9, velocity_bias_mode="off", velocity_bias_strength=0.0, velocity_bias_source="mean", velocity_refinement_key="default", rescue_mode=False, rescue_threshold=0.15, rescue_strength=0.2, rescue_prompt_sig=None, alg_enabled=False, alg_strength=2.0, alg_sigma_threshold=0.975, 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, h3_audio_clock=False): """ ODE sampler for distilled few-step video models (e.g. LTX2.3 distilled LoRA). - Adams-Bashforth 2-step multistep (order=2): extrapolates the denoised direction from two consecutive steps for second-order accuracy at zero extra model-call cost. Reduces discretisation error across the large sigma jumps typical of 4–8 step distilled schedules. - Heun predictor-corrector on final steps: calls the model a second time at sigma_next to correct the update direction. Significantly improves sharpness and detail in the steps that define the final output. - Optional s_noise: tiny ancestral-style noise injection for diversity. Default 0 = fully deterministic ODE (recommended for distilled models). - Velocity-bias capture/apply and reactive rescue around each model eval — the same memory machinery the Hybrid/RF samplers use. On few-step schedules the velocity targets may only match a step or two, so apply/rescue fire less often than on an 8-step run; they no-op cleanly when no target matches. - EXPERIMENTAL alg_enabled: Adaptive Low-Pass Guidance (arXiv:2506.08456) for i2v anchored chunks. Blurs the anchor frame the model sees while sigma is above alg_sigma_threshold (the near-pure-noise steps), then swaps back to the sharp anchor for the rest of the schedule — counters the model's tendency to shortcut to a near-static video that just matches the sharp reference. No-op when there's no i2v anchor (denoise_mask is None) or the packed latent layout can't be read. - EXPERIMENTAL alg_guide_tail_frames (>0 extends the same idea, not from the paper): additionally blurs that many trailing video frames — the newly-appended guide- attention frames this scene (mid_scene_guide / carry_i2v_guides-as-guide / configured guides / JoyAI memory), set by the Scene Chain Sampler's alg_blur_guides toggle right before this chunk's sample call. 0 = untouched (default). Standalone: works with alg_enabled off (anchor stays sharp, only the guide tail is blurred), with its own independent controls — alg_guide_blur_strength / alg_guide_blur_sigma_threshold (also set per-chunk by the Scene Chain Sampler), separate from the anchor's alg_strength / alg_sigma_threshold. - EXPERIMENTAL mg_enabled: Momentum Guidance (arXiv:2602.20360). Keeps a running EMA (decay=mg_decay) of the per-step ODE direction and blends the current step's direction toward it (weight=mg_strength) while sigma is BELOW mg_sigma_threshold — the complementary window to ALG's blur, which only acts above its threshold. The EMA itself accumulates every step (free) so it has real history by the time it starts being applied; only the application is gated. Audio is never touched (video-only). - Optional quality_sharpness: temporal-average unsharp on the x0 prediction, applied only during the final Heun-correction steps (same mechanism as the Hybrid Euler 2S sampler's quality-phase sharpening, ported here). Free (no extra model eval), video-only. - Optional h3_audio_clock: on MiniMax H3, integrate the audio stream on its own flow schedule instead of the tangent approximation the DiT hands the sampler. Matters most here — this sampler exists for the few-step schedules where that tangent is worst. Free (a scalar per step), no-op off H3. """ extra_args = {} if extra_args is None else extra_args seed = extra_args.get("seed", None) noise_sampler = k_diffusion_sampling.default_noise_sampler(x, seed=seed) total_steps = max(0, len(sigmas) - 1) if total_steps <= 0: return x sched_steps = max(1, int(len(sigmas)) - 1) order = max(1, min(2, int(order))) s_noise = max(0.0, min(1.0, float(s_noise))) # The TERMINAL step (sigma_next == 0) cannot be Heun-corrected: the corrector evaluates # the model at sigma_next, and sigma 0 is degenerate. It also returns early below, so a # window measured from the end of the schedule started ON that step and the corrector was # never reached — `final_correction_steps=1` did nothing at all on any schedule ending at # 0, which is all of them. The window is measured over the CORRECTABLE steps instead, so # N means N. Sharing the index with quality_sharpness is deliberate: its own tooltip # defines its window as the Heun-correction steps. correctable = sum(1 for i in range(sched_steps) if float(sigmas[i + 1]) > 0) final_correction_steps = max(0, min(correctable, int(final_correction_steps))) correction_start_idx = correctable - final_correction_steps _RESCUE_LOG["warned_no_memory"] = False _RESCUE_LOG["warned_no_prompt_match"] = False _RESCUE_LOG["fired"] = 0 # Audio-safe steering: on a packed LTXAV latent keep velocity bias / rescue on the # video stream only (perturbing the audio stream corrupts it via joint attention). # None for single-stream LTXV -> _video_only is a no-op. video_mask = _packed_video_mask(model, x) if video_mask is not None and (rescue_mode or velocity_bias_mode != "off"): n_aud = int((video_mask < 0.5).sum().item()) print(f"[FunPack AV] packed audio+video latent detected -> audio-safe steering " f"(audio held deterministic on {n_aud} of {video_mask.shape[-1]} packed dims)") normalize_strength = max(0.0, min(1.0, float(normalize_strength))) if normalize_strength > 0.0: print(f"[FunPack AV] latent normalization on (strength={normalize_strength}, " f"start_sigma={normalize_start_sigma}, video-only)") audio_clock = _audio_clock_setup(model, x, sigmas, h3_audio_clock) s_in = x.new_ones([x.shape[0]]) prev_denoised = None prev_h = None ref_scale = [None] # video-only latent-normalization reference (anti-overbake), opt-in _kvlock_state = {} # EXPERIMENTAL ALG (arXiv:2506.08456): blur the i2v anchor for the earliest/near-pure- # noise steps, then swap back to the sharp anchor. No-op without a real anchor (no # denoise_mask) or if the packed-latent layout can't be read (helper returns None). # alg_guide_tail_frames (>0) extends the same idea, beyond the paper, to the trailing # guide-attention frames this scene appended (alg_blur_guides on the Scene Chain Sampler). # Anchor and guide-tail blur are fully independent: each has its own strength and sigma # window (alg_strength/alg_sigma_threshold vs alg_guide_blur_*), so one precomputed # latent per (anchor blurred?, tail blurred?) combination that can occur — the frame # sets are disjoint, so the both-blurred variant just composes the two. alg_guide_tail_frames = max(0, int(alg_guide_tail_frames)) alg_sharp_latent_image, alg_latents, alg_anchor_on, alg_tail_on = _alg_prepare( model, extra_args, bool(alg_enabled), alg_strength, alg_guide_tail_frames, alg_guide_blur_strength, ) alg_active = alg_anchor_on or alg_tail_on if alg_active: anchor_desc = (f"strength={alg_strength}, sigma_threshold={alg_sigma_threshold}" if alg_anchor_on else "off") tail_desc = (f"{alg_guide_tail_frames} frame(s), strength={alg_guide_blur_strength}, " f"sigma_threshold={alg_guide_blur_sigma_threshold}" if alg_tail_on else "off") print(f"[FunPack AV] ALG on (anchor: {anchor_desc}; guide tail: {tail_desc}) " f"— blurred while sigma > threshold") # EXPERIMENTAL Momentum Guidance (arXiv:2602.20360): EMA of the per-step ODE direction, # blended into the direction actually used once sigma drops below mg_sigma_threshold — # the complementary window to ALG's blur. EMA accumulates every step regardless of # whether it's being applied yet, so there's no cold-start when it switches on. mg_active = bool(mg_enabled) mg_decay = max(0.0, min(0.999, float(mg_decay))) mg_strength = max(0.0, min(1.0, float(mg_strength))) mg_ema = None if mg_active: print(f"[FunPack AV] Momentum Guidance on (strength={mg_strength}, decay={mg_decay}, " f"sigma_threshold={mg_sigma_threshold}) — applied while sigma < threshold") def _mg_step(d, video_mask): nonlocal mg_ema if not mg_active: return d d_detached = d.detach() mg_ema = d_detached if mg_ema is None else (mg_decay * mg_ema + (1.0 - mg_decay) * d_detached) if float(sigma) >= float(mg_sigma_threshold) or mg_strength <= 0.0: return d try: blended = d * (1.0 - mg_strength) + mg_ema.to(device=d.device, dtype=d.dtype) * mg_strength except Exception as _e: _log.failed("FunPackSceneChain", "momentum guidance", _e, "this step uses the raw derivative, unsmoothed") return d return _video_only(blended, d, video_mask) try: 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) if alg_active: model.latent_image = alg_latents[( alg_anchor_on and float(sigma) > float(alg_sigma_threshold), alg_tail_on and float(sigma) > float(alg_guide_blur_sigma_threshold), )] 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), prompt_sig=rescue_prompt_sig, source=velocity_bias_source) x = _video_only(x, x_pre, video_mask) denoised = model(x, sigma * s_in, **extra_args) _kvlock_schedule(model, denoised, prev_denoised, video_mask, _kvlock_state) if rescue_mode or _velocity_bias_enabled(velocity_bias_mode, "capture"): _capture_velocity_bias(velocity_refinement_key, velocity_target, x, sigma, denoised, prompt_sig=rescue_prompt_sig) if rescue_mode and velocity_target is not None: denoised = _video_only(_rescue_denoised( denoised, x, sigma, velocity_refinement_key, velocity_target, rescue_threshold, rescue_strength, prompt_sig=rescue_prompt_sig, source=velocity_bias_source, ), denoised, video_mask) # 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: 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. denoised = _normalize_video_denoised( denoised, video_mask, sigma, ref_scale, normalize_strength, normalize_start_sigma, ) if callback is not None: callback({"x": x, "i": i, "sigma": sigma, "sigma_hat": sigma, "denoised": denoised}) h = float((sigma - sigma_next).abs().item()) # Adams-Bashforth 2-step multistep correction. # Coefficients for variable step sizes: r = h_current / h_previous. # denoised_eff = (1 + r/2) * denoised - (r/2) * prev_denoised if order >= 2 and prev_denoised is not None and prev_h is not None and prev_h > 1e-7 and h > 1e-7: r = max(0.1, min(5.0, h / prev_h)) try: denoised_eff = (1.0 + r / 2.0) * denoised - (r / 2.0) * prev_denoised.to(device=denoised.device, dtype=denoised.dtype) except Exception: denoised_eff = denoised else: denoised_eff = denoised # Graduated 2nd order (opt-in, free): ramp the AB2 contribution linearly 0->1 across # 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 sched_steps > 1: w = i / (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, # raw denoised for audio (2nd-order extrapolation corrupts the audio stream). denoised_eff = _video_only(denoised_eff, denoised, video_mask) # Store current denoised for the next step's multistep correction. # Reset after a Heun step since x was updated with a corrected direction. prev_denoised = denoised.detach() prev_h = h if sigma_next == 0: # Same displacement the euler branch below would produce at dt = -sigma, # so the audio clock corrects this final step like any other. x = _audio_clock_step(denoised_eff, x, audio_clock, i) continue dt = sigma_next - sigma # negative: sigmas decrease if i >= correction_start_idx: # Heun predictor-corrector. # Predictor: Euler step using the (multistep-corrected) denoised. d1 = k_diffusion_sampling.to_d(x, sigma, denoised_eff) # The corrector's own audio output is discarded (audio rides d1), but the # state it is evaluated at is shared: audio and video are one packed # sequence, so an over-stepped audio half would skew the VIDEO correction # through joint attention. Predict the audio where the clock says it lands. x_pred = _audio_clock_step(x + d1 * dt, x, audio_clock, i) # Corrector: evaluate model at the predicted x and sigma_next. denoised_pred = model(x_pred, sigma_next * s_in, **extra_args) d2 = k_diffusion_sampling.to_d(x_pred, sigma_next, denoised_pred) # Audio rides the plain euler direction (d1); only video gets the Heun correction. d_use = _video_only((d1 + d2) / 2.0, d1, video_mask) d_use = _mg_step(d_use, video_mask) x = _audio_clock_step(x + d_use * dt, x, audio_clock, i) # Heun updates x differently; invalidate multistep history. prev_denoised = None prev_h = None else: d = k_diffusion_sampling.to_d(x, sigma, denoised_eff) d = _mg_step(d, video_mask) x_det = _audio_clock_step(x + d * dt, x, audio_clock, i) if s_noise > 0.0 and float(sigma_next) > 0 and _rf_ancestral(model): # Rectified-flow ancestral step, identical to comfy's # sample_euler_ancestral_RF (which is what `euler_ancestral` becomes on a # CONST model such as MiniMax H3 or LTXAV). # # The previous version stepped the FULL distance to sigma_next and then # added sqrt(sigma^2 - sigma_next^2) of noise on top. Both halves are wrong # for a flow model: that variance formula is the VP one, and adding noise # without shortening the deterministic step leaves the latent noisier than # the schedule says it should be at sigma_next. It over-noised, which is # why s_noise here never behaved like euler_ancestral. # # The correct step lands SHORT (sigma_down), then renoises with the alpha # rescaling flow matching requires. At s_noise=1.0 this is euler_ancestral. downstep_ratio = 1 + (sigma_next / sigma - 1) * s_noise sigma_down = sigma_next * downstep_ratio alpha_ip1 = 1 - sigma_next alpha_down = 1 - sigma_down ratio = sigma_down / sigma x_anc = ratio * x + (1 - ratio) * denoised_eff renoise = (sigma_next ** 2 - sigma_down ** 2 * alpha_ip1 ** 2 / alpha_down ** 2) ** 0.5 x_anc = ((alpha_ip1 / alpha_down) * x_anc + noise_sampler(sigma, sigma_next) * renoise) # Video only: ancestral noise on the audio stream corrupts it, and the # audio half of x_det already carries the clock correction. x = _video_only(x_anc, x_det, video_mask) elif s_noise > 0.0 and float(sigma_next) > 0: # Non-CONST (eps/VP) model: the original formulation, unchanged. sigma_up = math.sqrt(max(0.0, float(sigma.item()) ** 2 - float(sigma_next.item()) ** 2)) x = _video_only(x_det + noise_sampler(sigma, sigma_next) * s_noise * sigma_up, x_det, video_mask) if sigma_up > 0.0 else x_det else: x = x_det finally: if alg_active: model.latent_image = alg_sharp_latent_image return x # FunPack sampler functions that accept h3_audio_clock. The hybrid entry point is here # because on a CONST model (which MiniMax H3 is) it delegates to _sample_const_rf_full, # which carries the correction; its own eps loop is unreachable on H3. _AUDIO_CLOCK_SAMPLERS = (sample_funpack_distilled_flow, sample_funpack_hybrid_euler_2s) class FunPackDistilledFlowSampler: @classmethod def INPUT_TYPES(cls): return { "required": { "order": ("INT", { "default": 2, "min": 1, "max": 2, "step": 1, "tooltip": "Multistep order. 1 = standard Euler ODE. 2 = Adams-Bashforth 2-step: extrapolates the denoised direction from two consecutive steps for better accuracy at no extra model-call cost.", }), "final_correction_steps": ("INT", { "default": 0, "min": 0, "max": 3, "step": 1, "tooltip": "Number of correctable steps that use a Heun predictor-corrector pass. EACH COSTS ONE EXTRA MODEL CALL (on a 7-step run, 1 = ~14% more time). Counted over the steps that CAN be corrected: the terminal step lands on sigma 0, where the corrector has nothing to evaluate, so it is never one of them. Was 1 by default and did nothing — the window used to start on that terminal step, which returns before the corrector runs — so 0 is what every run has actually been doing. Set it to 1 to get the correction the knob always advertised, on the last real step. quality_sharpness shares this window and needs this above 0.", }), "s_noise": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.00, "step": 0.01, "tooltip": "Ancestral noise: the fraction of each step that is taken stochastically rather than deterministically (comfy calls this eta). 0 = fully deterministic ODE. 1.0 on a rectified-flow model (MiniMax H3, LTXAV) is exactly euler_ancestral. Free — no extra model calls. BEHAVIOUR CHANGED: this used to step the full distance and then add sqrt(sigma^2-sigma_next^2) of noise ON TOP, which is the VP variance formula and leaves the latent noisier than the schedule says — it never matched euler_ancestral at any value. It now lands short (sigma_down) and renoises with the alpha rescaling flow matching needs. A value carried over from before does NOT mean what it used to: re-tune it. If you liked euler_ancestral, start at 1.0.", }), "velocity_bias_mode": (VELOCITY_BIAS_MODES, { "default": "off", "tooltip": "Experimental: capture/apply averaged early model velocity around normalized sigma 0.9/0.72/0.42. Off preserves the plain distilled ODE. Note: few-step schedules may only land on a target or two, so it fires less often than on an 8-step run.", }), "velocity_bias_strength": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 3.0, "step": 0.05, "tooltip": "Strength of the remembered velocity (action) injected at the structure sigma. 0 disables. ~0.15 = subtle spice; 0.3-1.0 = clear action crossover; 2-3 approaches full action replacement (capped so the current gen isn't wiped). Creative tool, not for consistency.", }), "velocity_bias_source": (["mean", "nearest"], { "default": "mean", "tooltip": "How velocity bias / rescue pick a good direction. 'mean' = prompt-blind global average. 'nearest' = single best-matching prompt cluster — preserves one real good gen's detail instead of a washed-out average. Affects both apply and rescue.", }), "velocity_refinement_key": ("STRING", { "default": "default", "multiline": False, "tooltip": "Memory key used to capture/apply early velocity bias and rescue trajectories.", }), "rescue_mode": ("BOOLEAN", { "default": False, "tooltip": "Reactive in-flight rescue, rating-gated. Steers each eligible step toward trajectories you rated good and away from ones you rated Awful (matched to the current prompt). A no-op until you've rated a few gens for this prompt/key.", }), "rescue_threshold": ("FLOAT", { "default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Fires when the step has diverged from the good trajectory by more than this (1 - cosine) OR aligned with a bad trajectory by more than this (cosine). Lower = corrects more eagerly. 0.10-0.20 typical.", }), "rescue_strength": ("FLOAT", { "default": 0.2, "min": 0.0, "max": 0.5, "step": 0.01, "tooltip": "How hard to pull toward good / push away from bad when triggered (magnitude preserved, no energy injected). Keep moderate; 0.5 is a strong correction.", }), }, "optional": { "sigmas": ("SIGMAS",), # Appended last on purpose: a new widget inserted earlier shifts every saved # widgets_values slot in existing Distilled Flow nodes. "ab2_ramp": ("BOOLEAN", { "default": False, "tooltip": "Graduated 2nd order (free). Instead of full AB2 on every step, ramp the AB2 contribution linearly 0->1 across the schedule: early/noisy steps stay near 1st-order euler (less overshoot), late/detail steps get full AB2. No extra model calls. Helps low-step distilled runs. No effect at order=1.", }), "normalize_strength": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Video-only latent normalization (anti-overbake / oversaturation / colour drift) stacked on this ODE. 0 = off. 0.5 = gentle. Audio is never touched. ~zero overhead.", }), "normalize_start_sigma": ("FLOAT", { "default": 0.9, "min": 0.0, "max": 1.0, "step": 0.025, "tooltip": "Sigma at/below which latent normalization activates and anchors its reference (above it the x0 estimate is meaningless). Only used when normalize_strength > 0.", }), # Appended last on purpose, same rule as ab2_ramp above. "alg_enabled": ("BOOLEAN", { "default": False, "tooltip": "EXPERIMENTAL: Adaptive Low-Pass Guidance (arXiv:2506.08456). Blurs the i2v anchor frame while sigma is above alg_sigma_threshold, then swaps back to the sharp anchor — counters the model's tendency to shortcut to a near-static video that just matches the reference image. No-op without an i2v anchor.", }), "alg_strength": ("FLOAT", { "default": 2.0, "min": 1.0, "max": 4.0, "step": 0.1, "tooltip": "Downsample factor for the anchor blur (paper default 2.5, but 2.0 held character/i2v consistency noticeably better in testing here). Higher = blurrier anchor during the affected steps. Only used when alg_enabled; guide-frame blur has its own controls on the Scene Chain Sampler (alg_blur_guides + alg_guide_blur_*).", }), "alg_sigma_threshold": ("FLOAT", { "default": 0.975, "min": 0.5, "max": 0.999, "step": 0.005, "tooltip": "Anchor stays blurred while sigma is above this value (the near-pure-noise steps), then swaps to sharp. Higher = narrower blurred window. Only used when alg_enabled; guide-frame blur has its own controls on the Scene Chain Sampler (alg_blur_guides + alg_guide_blur_*).", }), "mg_enabled": ("BOOLEAN", { "default": False, "tooltip": "EXPERIMENTAL: Momentum Guidance (arXiv:2602.20360). Keeps a running average of the per-step direction and blends the current step toward it once sigma drops below mg_sigma_threshold — the complementary window to ALG's blur. Smooths the fine-motion/refinement steps; may damp motion as a side effect, untested for video.", }), "mg_strength": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Blend weight toward the momentum average when active (0 = no effect, 1 = fully replace the step's direction with the average). Only used when mg_enabled.", }), "mg_decay": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 0.99, "step": 0.05, "tooltip": "EMA decay for the momentum average (higher = longer memory of past steps' directions). On our 8-step schedule, high decay (e.g. 0.9) keeps the EMA anchored to the very first, near-pure-noise step's direction for nearly the whole run — wrong scale, produces garbage regardless of mg_strength. Tested safe (and good) at 0.5 even at mg_strength=1.0. Only used when mg_enabled.", }), "mg_sigma_threshold": ("FLOAT", { "default": 0.975, "min": 0.5, "max": 0.999, "step": 0.005, "tooltip": "Momentum guidance applies while sigma is BELOW this value (the opposite window from alg_sigma_threshold) — defaults to the same boundary as ALG for a clean handoff. Only used when mg_enabled.", }), # Appended last on purpose, same rule as ab2_ramp above. "quality_sharpness": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Restores fine detail via temporal-average unsharp on the x0 prediction, applied only during the final Heun-correction steps (final_correction_steps > 0 required to have any effect). 0 disables. 0.2-0.4 typical. Free (no extra model eval), video-only.", }), } } RETURN_TYPES = ("SAMPLER", "SIGMAS") RETURN_NAMES = ("sampler", "sigmas") FUNCTION = "get_sampler" CATEGORY = "FunPack/Sampling" DESCRIPTION = ( "ODE sampler for distilled few-step video models (e.g. LTX2.3 distilled LoRA). " "Adams-Bashforth 2-step multistep for better trajectory accuracy across large sigma jumps, " "Heun predictor-corrector on final steps for quality, optional controlled noise for diversity, " "optional velocity bias + reactive rescue (shared with the Hybrid/RF samplers), " "experimental ALG anchor de-staticking for i2v chunks, and experimental Momentum " "Guidance smoothing for the complementary (fine-motion) sigma window." ) def get_sampler(self, order=2, final_correction_steps=0, s_noise=0.0, velocity_bias_mode="off", velocity_bias_strength=0.0, velocity_bias_source="mean", velocity_refinement_key="default", rescue_mode=False, rescue_threshold=0.15, rescue_strength=0.2, rescue_prompt_sig=None, sigmas=None, ab2_ramp=False, normalize_strength=0.0, normalize_start_sigma=0.9, alg_enabled=False, alg_strength=2.0, alg_sigma_threshold=0.975, mg_enabled=False, mg_strength=0.5, mg_decay=0.5, mg_sigma_threshold=0.975, quality_sharpness=0.0): prepared_sigmas = sigmas.detach().clone() if isinstance(sigmas, torch.Tensor) else sigmas sampler = comfy.samplers.KSAMPLER( sample_funpack_distilled_flow, extra_options={ "order": order, "final_correction_steps": final_correction_steps, "ab2_ramp": ab2_ramp, "normalize_strength": normalize_strength, "normalize_start_sigma": normalize_start_sigma, "s_noise": s_noise, "velocity_bias_mode": velocity_bias_mode, "velocity_bias_strength": velocity_bias_strength, "velocity_bias_source": velocity_bias_source, "velocity_refinement_key": velocity_refinement_key, "rescue_mode": rescue_mode, "rescue_threshold": rescue_threshold, "rescue_strength": rescue_strength, "rescue_prompt_sig": rescue_prompt_sig, "alg_enabled": alg_enabled, "alg_strength": alg_strength, "alg_sigma_threshold": alg_sigma_threshold, "mg_enabled": mg_enabled, "mg_strength": mg_strength, "mg_decay": mg_decay, "mg_sigma_threshold": mg_sigma_threshold, "quality_sharpness": quality_sharpness, } ) return (sampler, prepared_sigmas) def _normalize_video_denoised(denoised, video_mask, sigma, ref, strength, start_sigma, tolerance=0.05, beta=0.9): """Keep the VIDEO latent's spread from drifting into overbaked / oversaturated ranges. Operates on the denoised (x0) prediction, video stream only — audio is never touched, so it is audio-safe by construction. Once x0 is meaningful (sigma <= start_sigma) it anchors a reference spread (robust std over the video region) and gently compresses any later step whose spread inflates past it; an EMA lets the reference follow legitimate detail growth while resisting abrupt overbake spikes. Benefit-only: it never amplifies and is a no-op when the latent stays healthy. Cost is a few masked reductions per step (~zero overhead). `ref` is a 1-element list holding the reference scale across steps (mutable state). """ if strength <= 0.0: return denoised try: s = float(sigma.item()) if hasattr(sigma, "item") else float(sigma) except Exception: s = 1.0 if s > start_sigma: return denoised # skip the pure-noise plateau where the x0 estimate is not yet meaningful try: # Stats over the video region (whole tensor for single-stream LTXV, where mask is None). if video_mask is not None: m = video_mask w = m.sum().clamp_min(1.0) mean = (denoised * m).sum() / w var = (((denoised - mean) ** 2) * m).sum() / w else: mean = denoised.mean() var = ((denoised - mean) ** 2).mean() std_v = float(torch.sqrt(var.clamp_min(1e-12)).item()) if ref[0] is None: ref[0] = std_v return denoised target = ref[0] * (1.0 + tolerance) if std_v > target and std_v > 1e-6: factor = 1.0 - strength * (1.0 - target / std_v) # = lerp(1, target/std, strength) normalized = mean + (denoised - mean) * factor denoised = _video_only(normalized, denoised, m) ref[0] = ref[0] * beta + target * (1.0 - beta) else: ref[0] = ref[0] * beta + std_v * (1.0 - beta) except Exception: return denoised return denoised class FunPackLTXAVSceneChainSampler: @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "vae": ("VAE",), "positive": ("CONDITIONING",), "negative": ("CONDITIONING",), "sampler": ("SAMPLER",), "sigmas": ("SIGMAS",), "seed": ("INT", {"default": 1, "min": 0, "max": 0xffffffffffffffff}), "latent_template": ("LATENT",), "num_frames_per_scene": ("INT", {"default": 97, "min": 1, "max": 4096, "step": 8}), "frame_overlap": ("INT", { "default": 16, "min": 0, "max": 512, "step": 8, "tooltip": "Pixel frames copied from the previous scene into the next chunk and preserved during denoising. 0 disables overlap blending entirely. WARNING: combining frame_overlap=0 with carry_i2v_guides=True is confirmed to produce bad results — use only for testing.", }), "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1}), "max_scenes": ("INT", {"default": 8, "min": 1, "step": 1}), "use_same_seed": ("BOOLEAN", { "default": False, "tooltip": "Use the first provided scene seed for every scene. Off uses per-scene metadata seeds or seed + scene index.", }), "carry_i2v_guides": ("BOOLEAN", { "default": False, "tooltip": "Carry protected frames from latent_template noise_mask into each continuation chunk as a style guide.", }), "mid_scene_guide": ("BOOLEAN", { "default": False, "tooltip": "Experimental: append the middle frame of the previous scene as a guide for the current scene via LTX guide attention. Helps maintain character positioning across scenes.", }), "mid_scene_guide_strength": ("FLOAT", { "default": 0.25, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Guide attention strength for mid-scene anchor. Full 0..1 range. 0.25-0.35 is the measured sweet spot: below 0.25 audio degrades and character appearance drifts, above 0.35 spatial conflicts appear when scene composition shifts. Outside that band is yours to explore — 0 disables the guide's pull entirely.", }), "embed_guidance": ("BOOLEAN", { "default": False, "tooltip": "Apply the Refiner's learned quality direction at each denoising step, not just once before sampling. Requires refinement_key_input and enough liked generations to have a direction. Adds ~20-30% inference overhead.", }), "embed_guidance_strength": ("FLOAT", { "default": 0.02, "min": 0.005, "max": 0.1, "step": 0.005, "tooltip": "Per-step nudge strength toward the liked conditioning direction. Keep small — the direction is applied at every step so it compounds. 0.01-0.03 is typical.", }), "score_slider": ("BOOLEAN", { "default": False, "tooltip": "FreeSliders-style taste guidance in SCORE space. Instead of nudging the conditioning once (embed_guidance), it runs 3 forward passes on quality-phase steps — base, taste+, taste- — and steers the noise prediction along eps_+ minus eps_-. Stronger, prompt-faithful taste push; ~2x cost on late steps. Uses the same learned direction + source + refinement_key_input as embed_guidance (needs 3+ liked generations). Contrastive pair: once 3+ disliked/awful gens are rated, the minus pole switches from a mirror of liked to the real learned BAD direction, so the axis becomes good-vs-bad and actively steers away from what produced rated-bad gens. Video-only (audio unaffected).", }), "score_slider_strength": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 3.0, "step": 0.25, "tooltip": "Slider amount (eta). How hard to push the noise prediction along the learned taste axis. 1.0 is a clear, safe push; raise toward 3.0 for a stronger effect (paper's saturation range). 0 = off.", }), "transition_duration": ("INT", { "default": 16, "min": 0, "max": 128, "step": 2, "tooltip": "Extra pixel frames of fade beyond the blend zone on each side of a scene boundary. 0 = disable all transition effects.", }), }, "optional": { # Widget inputs first, then connection/forceInput sockets LAST. A forceInput # input placed between widgets desyncs ComfyUI's widgets_values mapping on reload # (the combo value lands on decode_tile_size -> NaN/"relative" in the INT field). "decode_tile_size": ("INT", { "default": 0, "min": 0, "max": 4096, "step": 64, "tooltip": "Tile size for VAE decode (0 = no tiling). Set to e.g. 512 if decode OOMs.", }), "decode_noise_scale": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.005, "tooltip": "Noise injected at VAE decode to restore fine detail/grain (the LTX VAE decoder is itself a tiny diffusion model). 0 = off (clean decode). ~0.025 is a gentle detail restore. Applied to this node's IMAGES decode only.", }), "decode_timestep": ("FLOAT", { "default": 0.05, "min": 0.0, "max": 1.0, "step": 0.001, "tooltip": "Decode timestep fed to the VAE decoder when decode_noise_scale > 0. ~0.05 adds subtle detail; higher gives the decoder more freedom (more deviation from the latent).", }), "embed_guidance_source": (["relative", "absolute"], { "default": "relative", "tooltip": "Which learned direction embed_guidance steers toward. Relative: this prompt's liked direction (needs refinement_key_input). Absolute: the global, prompt-agnostic taste direction the Refiner accumulates across all prompts — works with no key.", }), # NOTE: append-only. New widgets MUST go at the END of the widget sequence (after # every widget the reference workflow already has a positional value for, before the # forceInput sockets below). Inserting earlier desyncs extract_widgets' positional # mapping and lands a numeric on a combo (e.g. joyai_frame_select got 0). "joyai_memory": ("BOOLEAN", { "default": False, "tooltip": "JoyAI-Echo cross-shot memory bank. REQUIRES the JoyAI-Echo LoRA — without it the injected memory frames change nothing, because the base model was never trained to read them as memory. Generalizes mid_scene_guide from one anchor to a managed set of clean prior-shot frames injected into each scene via LTX guide attention, so character/scene identity carries across the whole chain (JoyAI-Echo's story-level consistency). The first joyai_fix_frames scenes are pinned permanently as a global anchor; the rest is a rolling most-recent window capped at joyai_memory_size. Supersedes mid_scene_guide when on. Video memory only; pair it with joyai_audio_memory for the soundtrack.", }), "joyai_memory_size": ("INT", { "default": 7, "min": 1, "max": 32, "tooltip": "Max total memory entries injected per scene (JoyAI default 7). Higher = stronger long-range consistency but more guide tokens and slower scenes.", }), "joyai_fix_frames": ("INT", { "default": 3, "min": 0, "max": 16, "tooltip": "Number of opening scenes pinned permanently in the bank as a global anchor (JoyAI default 3). They are never pruned; entries beyond them are a rolling most-recent window.", }), "joyai_frame_select": (["center", "first", "random"], { "default": "center", "tooltip": "Which frame of each finished scene to store in the bank (JoyAI default 'center').", }), "joyai_memory_strength": ("FLOAT", { "default": 0.3, "min": 0.25, "max": 10.0, "step": 0.05, "tooltip": "Guide-attention strength for each memory frame. 0.25 floor as mid_scene_guide (below it audio degrades and identity drifts). Uncapped at the top: 0.25-0.5 is the audio-safe band, higher values push identity harder but may degrade audio/over-constrain motion.", }), # NOTE: append-only — keep new sampler widgets at the END of this block so the # builder's positional reference-workflow mapping (extract_widgets) stays aligned. "joyai_audio_memory": ("BOOLEAN", { "default": False, "tooltip": "JoyAI-Echo PAIRED AUDIO memory. Needs the JoyAI-Echo LoRA, like the video memory it accompanies. Alongside each video memory frame, pin the prior shot's clean audio latent into the audio stream so voice/timbre/ambience carry across shots the way the face now does. Deliberately breaks the audio pass-through invariant — off by default. Requires joyai_memory on; no effect on single-stream (video-only) LTXV.", }), "v2a_grad_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 4.0, "step": 0.25, "tooltip": "JoyAI-Echo video->audio coupling. Scales the model's trained video-to-audio cross-attention so the carried audio tracks the new shot's visuals (JoyAI uses 2.0). 1.0 = native model behavior (no change, zero overhead); 0.0 = audio ignores video this run. Only applies when joyai_audio_memory is on.", }), "audio_vae": ("VAE", { "tooltip": "MiniMax H3 only: the audio VAE, needed to encode AUDIO reference media for ref2va (a voice or ambience clip the generation should sound like). Studio lists the references and bakes their