From ce2de22b2eba002b674b1c4ce194c79e6ce894bf Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Mon, 28 Sep 2026 19:35:29 +0300 Subject: [PATCH] feat: STAS for H3; late-branch guidance skips negative-prompt calls STAS (arXiv 2603.17825) sets massive-activation channels on frame 0 and frame edges to alpha x peak at one block over the first 40% of steps. Rating-learned dials now share rated_dial.py. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 + conditioning.py | 13 +- late_guidance.py | 95 +----------- movie_editor/backend/pipeline_caps.py | 1 + movie_editor/backend/settings_card.py | 5 + movie_editor/frontend/engine_settings.js | 10 +- movie_editor/frontend/pipeline_caps.js | 5 +- rated_dial.py | 90 +++++++++++ samplers.py | 188 ++++++++++++++++++++--- stas.py | 66 ++++++++ tests/test_late_guidance.py | 43 ++++-- tests/test_stas.py | 97 ++++++++++++ 12 files changed, 483 insertions(+), 132 deletions(-) create mode 100644 rated_dial.py create mode 100644 stas.py create mode 100644 tests/test_stas.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5e08603f..8332a3ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ varies more. Learned from ratings or set by hand. Picture only, no extra model call. - **Late-branch guidance** (experimental, H3): each step also makes a weaker picture (one late block skipped) and steers away from it. ~15% cost at block 43. Learned or manual. +- **STAS** (experimental, H3): strengthens the model's own frame-edge signals over the first + steps, for steadier motion. No extra model call. Learned or manual. - **Coloured logs**: red errors, orange warnings, green successes, in the terminal (FunPack's lines; `FUNPACK_LOG_COLOR=1/0` forces) and in the Editor's Log panel. The panel now also shows ComfyUI's own logging output, node crashes included. diff --git a/conditioning.py b/conditioning.py index 707f97d0..02ff42d1 100644 --- a/conditioning.py +++ b/conditioning.py @@ -10507,21 +10507,24 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): else float(learning_profile.get("reward", 0.0) or 0.0)) try: try: - from . import shot_memory as _sm, tsr as _tsr_mod, late_guidance as _lg + from . import shot_memory as _sm, tsr as _tsr_mod, late_guidance as _lg, stas as _st except ImportError: - import shot_memory as _sm, tsr as _tsr_mod, late_guidance as _lg + import shot_memory as _sm, tsr as _tsr_mod, late_guidance as _lg, stas as _st _n = _sm.commit(refinement_key, _sm_reward) if _n is not None: print(f"[FunPackRefiner] shot memory: {_n} rated shot(s) on this key") _n = _tsr_mod.commit(refinement_key, _sm_reward) if _n is not None: print(f"[FunPackRefiner] decisiveness: {_n} rating(s) on this key") - _n = _lg.commit(refinement_key, _sm_reward) + _n = _lg.DIAL.commit(refinement_key, _sm_reward) if _n is not None: print(f"[FunPackRefiner] late-branch guidance: {_n} rating(s) on this key") + _n = _st.DIAL.commit(refinement_key, _sm_reward) + if _n is not None: + print(f"[FunPackRefiner] STAS: {_n} rating(s) on this key") except Exception as _e: - print(f"[FunPackRefiner] shot memory / decisiveness / late-branch guidance " - f"intake failed: {_e}") + print(f"[FunPackRefiner] shot memory / decisiveness / late-branch guidance / " + f"STAS intake failed: {_e}") # Absolute store: the same rating also feeds the keyless, prompt-agnostic taste prior. # Runs even with no refinement_key (Absolute is global), so standalone runs still build it. # Skipped for Wrong-* repair ratings (skip_value_function): Absolute reads reward as pure diff --git a/late_guidance.py b/late_guidance.py index 77e9a93a..1b6a6bc5 100644 --- a/late_guidance.py +++ b/late_guidance.py @@ -6,100 +6,17 @@ Branching at 43 of H3's 50 blocks costs ~7 blocks per step (~15%), not a second guided = normal + w * (normal - weak) picture only; sound is the normal pass -One number per run, learned from ratings the same way as decisiveness (tsr.py): each run -tries a strength near the key's current value, and each rating pulls the value toward the -strength that was liked or away from the one that was disliked. +The strength w is learned from ratings (rated_dial.py) or set by hand. """ -import json -import os -import random +try: + from .rated_dial import Dial, MODES # noqa: F401 +except ImportError: + from rated_dial import Dial, MODES # noqa: F401 -START = 0.5 -W_MAX = 1.5 -EXPLORE = 0.15 -LR = 0.5 -MODES = ("off", "learned", "manual") +DIAL = Dial("late_guidance", start=0.5, lo=0.0, hi=1.5, explore=0.15) DEFAULT_BLOCK = 43 -def _path(refinement_key): - try: - from .conditioning import refinement_state_path - except ImportError: - from conditioning import refinement_state_path - return refinement_state_path(refinement_key, "late_guidance", prefix="refine_v2", - extension="json") - - -def _read(refinement_key): - try: - with open(_path(refinement_key), "r", encoding="utf-8") as f: - data = json.load(f) - return data if isinstance(data, dict) else {} - except (FileNotFoundError, ValueError): - return {} - - -def _write(refinement_key, data): - path = _path(refinement_key) - os.makedirs(os.path.dirname(path), exist_ok=True) - tmp = path + ".tmp" - with open(tmp, "w", encoding="utf-8") as f: - json.dump(data, f) - os.replace(tmp, path) - - -def learned_w(history): - m = START - for h in history: - s = (h["reward"] > 0) - (h["reward"] < 0) - m += LR * s * (float(h["w"]) - m) - m = min(max(m, 0.0), W_MAX) - return m - - -def choose_w(refinement_key, mode, manual_w=START, rng=random): - """-> (w, note). A learned w only becomes pending once the run finishes (save_pending).""" - if mode == "manual": - return max(0.0, float(manual_w)), f"manual strength {float(manual_w):.2f}" - if mode != "learned" or not refinement_key: - return 0.0, "" - history = _read(refinement_key).get("history") or [] - centre = learned_w(history) - w = min(max(centre + rng.gauss(0.0, EXPLORE), 0.0), W_MAX) - return w, f"learned strength {centre:.2f} from {len(history)} rating(s), trying {w:.2f}" - - def mix(normal, weak, w): return normal + w * (normal - weak) - - -def save_pending(refinement_key, w): - if not refinement_key: - return - data = _read(refinement_key) - data["pending"] = float(w) - _write(refinement_key, data) - - -def clear_pending(refinement_key): - """A keyed run not learning w: its rating must not score an older run's w.""" - if refinement_key and "pending" in _read(refinement_key): - data = _read(refinement_key) - data.pop("pending", None) - _write(refinement_key, data) - - -def commit(refinement_key, reward): - """Pair the pending w with the rating of its run. -> ratings held, or None.""" - if not refinement_key: - return None - data = _read(refinement_key) - w = data.pop("pending", None) - if w is None: - return None - if reward: - data["history"] = ((data.get("history") or []) + [{"w": float(w), "reward": float(reward)}])[-200:] - _write(refinement_key, data) - return len(data["history"]) if reward else None diff --git a/movie_editor/backend/pipeline_caps.py b/movie_editor/backend/pipeline_caps.py index 54224a68..d75f5c01 100644 --- a/movie_editor/backend/pipeline_caps.py +++ b/movie_editor/backend/pipeline_caps.py @@ -126,6 +126,7 @@ SIMPLE_MODE_SAMPLER_OFF: dict[str, Any] = { "shot_memory": "off", "tsr": "off", "late_guidance": "off", + "stas": "off", } # Studio refiner keys, same rule. diff --git a/movie_editor/backend/settings_card.py b/movie_editor/backend/settings_card.py index 785d1977..ad516e4f 100644 --- a/movie_editor/backend/settings_card.py +++ b/movie_editor/backend/settings_card.py @@ -217,6 +217,8 @@ _OWNED_BY = { "tsr_k": ("tsr", ('manual',)), "late_guidance_strength": ("late_guidance", ('manual',)), "late_guidance_block": ("late_guidance", ('learned', 'manual')), + "stas_alpha": ("stas", ('manual',)), + "stas_block": ("stas", ('learned', 'manual')), "h3_block_repeat_span_loop": ("h3_block_repeat", None), "h3_block_repeat_video_only": ("h3_block_repeat", None), "h3_block_repeat_last_steps": ("h3_block_repeat", None), @@ -297,6 +299,9 @@ _EDITOR_DEFAULTS = { "late_guidance": "off", "late_guidance_strength": 0.5, "late_guidance_block": 43, + "stas": "off", + "stas_alpha": 2.0, + "stas_block": 15, "embed_guidance": False, "embed_guidance_source": "relative", "embed_guidance_strength": 0.02, diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index 62e592c8..20a086e6 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -25,6 +25,7 @@ "dynashift", "dynashift_strength", "dynashift_threshold", "shot_memory", "shot_memory_amount", "tsr", "tsr_k", "late_guidance", "late_guidance_strength", "late_guidance_block", + "stas", "stas_alpha", "stas_block", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block", ]); const RATING_GATED_STUDIO = new Set(["reference_injection", "value_guidance", "steer_mode", "absolute_strength"]); @@ -327,6 +328,13 @@ hint: "How hard to push away from the weak copy. 0 = off." }, { name: "late_guidance_block", label: "Late-branch block", kind: "int", default: 43, min: 1, max: 49, dependsOn: "late_guidance", dependsVals: ["learned", "manual"], hint: "The block the weak copy skips. Higher = cheaper and subtler." }, + { name: "stas", label: "STAS: steadier motion (H3, experimental)", kind: "combo", choices: ["off", "learned", "manual"], default: "off", + hint: "Strengthens the model's own frame-edge signals over the first steps. Aimed at steadier motion. Learned: ratings set the strength.", + detail: "Massive-activation steering (arXiv 2603.17825). A few channels carry huge values on the first frame and each frame's edges; this sets them to alpha times their peak at one block, over the first 40% of steps. No extra model call. Tested on Wan and CogVideoX, untested on H3; may calm subtle motion. The run log says which channels it found, or that it found none at that block. Learned needs a refinement key." }, + { name: "stas_alpha", label: "STAS strength", kind: "float", default: 2.0, min: 0.5, max: 4.0, step: 0.1, dependsOn: "stas", dependsValue: "manual", + hint: "Paper range 1.2-2.5." }, + { name: "stas_block", label: "STAS block", kind: "int", default: 15, min: 0, max: 49, dependsOn: "stas", dependsVals: ["learned", "manual"], + hint: "The block whose output is steered. The paper's best was ~30% deep." }, { name: "h3_repr_steering", label: "REINS (representation steering, H3, experimental)", kind: "bool", default: false, hint: "EXPERIMENTAL, unvalidated, H3 only. Reaches inside the model instead of around it: captures each steered block's video-row hidden state every generation, and once a block has 3+ liked and 3+ disliked runs of its OWN, adds that block's mean liked-minus-disliked difference back into its own output on every later run. No architectural ceiling like the attention-bias mechanisms have — push the strength too far and coherence can break with no warning." }, { name: "h3_repr_steering_strength", label: "REINS strength", kind: "float", default: 0.05, min: 0.0, max: 2.0, step: 0.01, dependsOn: "h3_repr_steering", @@ -550,7 +558,7 @@ const CHAIN_VIEW_KNOBS = { chain_continuity: ["carry_i2v_guides", "carry_overlap_through_anchor"], chain_timing: ["frame_overlap", "transition_duration", "use_same_seed", "cut_opening_frames"], - chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "taste_nearest_prompt", "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "explore_first_step", "explore_first_step_candidates", "dynashift", "dynashift_strength", "dynashift_threshold", "shot_memory", "shot_memory_amount", "tsr", "tsr_k", "late_guidance", "late_guidance_strength", "late_guidance_block", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block", "h3_repr_steering_passive_capture", "h3_q_steer_block", "h3_q_steer_strength", "h3_av_decouple"], + chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "taste_nearest_prompt", "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "explore_first_step", "explore_first_step_candidates", "dynashift", "dynashift_strength", "dynashift_threshold", "shot_memory", "shot_memory_amount", "tsr", "tsr_k", "late_guidance", "late_guidance_strength", "late_guidance_block", "stas", "stas_alpha", "stas_block", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block", "h3_repr_steering_passive_capture", "h3_q_steer_block", "h3_q_steer_strength", "h3_av_decouple"], chain_decode: ["h3_video_detail", "decode_noise_scale", "decode_timestep", "decode_tile_size"], chain_experimental: ["context_windows", "context_window_length", "context_window_overlap", "context_window_schedule", "context_window_fuse", "context_window_freenoise", "context_window_retain_first", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "alg_guide_blur_strength", "alg_guide_blur_sigma_threshold", "bounded_attention_enabled", "h3_block_repeat", "h3_block_repeat_span_loop", "h3_block_repeat_video_only", "h3_block_repeat_times", "h3_block_repeat_last_steps", "h3_explore_temperature", "h3_explore_temperature_block", "h3_shadow_negative", "h3_shadow_negative_compose", "h3_shadow_negative_video_scale", "h3_shadow_negative_audio_scale", "h3_shadow_negative_tau", "h3_shadow_negative_alpha", "h3_shadow_negative_start_percent", "h3_shadow_negative_end_percent", "identity_transfer_enabled", "source_id", "phase_scale", "id_strength", "arcface_mode", "debug_log"], }; diff --git a/movie_editor/frontend/pipeline_caps.js b/movie_editor/frontend/pipeline_caps.js index 58749aa4..c5438704 100644 --- a/movie_editor/frontend/pipeline_caps.js +++ b/movie_editor/frontend/pipeline_caps.js @@ -180,6 +180,9 @@ late_guidance: "Skips one of H3's own DiT blocks in a weak copy of each step. LTX runs a real CFG " + "negative, and its blocks are hooked differently.", + stas: + "Steers H3's own block outputs by their packed frame layout; LTX's blocks are hooked " + + "differently.", h3_video_detail: "Scales H3's own final layer, which reads video and audio on separate rows. LTX has no " + "layer of that shape, so the sampler has nothing to scale.", @@ -191,7 +194,7 @@ "h3_shadow_negative_tau", "h3_shadow_negative_alpha", "h3_shadow_negative_start_percent", "h3_shadow_negative_end_percent", "h3_shadow_negative_compose", - "late_guidance_strength", "late_guidance_block", + "late_guidance_strength", "late_guidance_block", "stas_alpha", "stas_block", ]; // Sub-settings of identity transfer: meaningless wherever the feature itself cannot run, diff --git a/rated_dial.py b/rated_dial.py new file mode 100644 index 00000000..e9ffdbc9 --- /dev/null +++ b/rated_dial.py @@ -0,0 +1,90 @@ +"""One number per run, learned from Liked/Disliked ratings. + +Each run tries a value near the key's current centre; the rating of that run pulls the +centre toward the tried value (liked) or pushes it away (disliked). The tried value is +only pending once the run finishes, so an interrupted run never takes a finished run's +rating, and a keyed run with the feature off clears it. Stored as +`..json`: {"history": [{"v", "reward"}], "pending": v}. +""" + +import json +import os +import random + +MODES = ("off", "learned", "manual") + + +class Dial: + def __init__(self, name, start, lo, hi, explore, lr=0.5): + self.name, self.start, self.lo, self.hi = name, start, lo, hi + self.explore, self.lr = explore, lr + + def path(self, refinement_key): + try: + from .conditioning import refinement_state_path + except ImportError: + from conditioning import refinement_state_path + return refinement_state_path(refinement_key, self.name, prefix="refine_v2", + extension="json") + + def read(self, refinement_key): + try: + with open(self.path(refinement_key), "r", encoding="utf-8") as f: + data = json.load(f) + return data if isinstance(data, dict) else {} + except (FileNotFoundError, ValueError): + return {} + + def _write(self, refinement_key, data): + path = self.path(refinement_key) + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump(data, f) + os.replace(tmp, path) + + def learned(self, history): + m = self.start + for h in history: + s = (h["reward"] > 0) - (h["reward"] < 0) + m += self.lr * s * (float(h["v"]) - m) + m = min(max(m, self.lo), self.hi) + return m + + def choose(self, refinement_key, mode, manual, rng=random): + """-> (value, note); `None` value = off this run.""" + if mode == "manual": + return float(manual), f"manual {float(manual):.2f}" + if mode != "learned" or not refinement_key: + return None, "" + history = self.read(refinement_key).get("history") or [] + centre = self.learned(history) + v = min(max(centre + rng.gauss(0.0, self.explore), self.lo), self.hi) + return v, f"learned {centre:.2f} from {len(history)} rating(s), trying {v:.2f}" + + def save_pending(self, refinement_key, v): + if refinement_key: + data = self.read(refinement_key) + data["pending"] = float(v) + self._write(refinement_key, data) + + def clear_pending(self, refinement_key): + """A keyed run not learning this: its rating must not score an older run's value.""" + if refinement_key and "pending" in self.read(refinement_key): + data = self.read(refinement_key) + data.pop("pending", None) + self._write(refinement_key, data) + + def commit(self, refinement_key, reward): + """Pair the pending value with its run's rating. -> ratings held, or None.""" + if not refinement_key: + return None + data = self.read(refinement_key) + v = data.pop("pending", None) + if v is None: + return None + if reward: + data["history"] = ((data.get("history") or []) + + [{"v": float(v), "reward": float(reward)}])[-200:] + self._write(refinement_key, data) + return len(data["history"]) if reward else None diff --git a/samplers.py b/samplers.py index 04f52990..35370b3d 100644 --- a/samplers.py +++ b/samplers.py @@ -22,11 +22,13 @@ try: from . import shot_memory as _shot_memory from . import tsr as _tsr from . import late_guidance as _late + from . import stas as _stas except ImportError: # flat import when ComfyUI loads the pack as a top-level module import funpack_log as _log import shot_memory as _shot_memory import tsr as _tsr import late_guidance as _late + import stas as _stas MOTION_PULSE_MODES = ["off", "balanced", "aggressive", "custom"] @@ -851,6 +853,22 @@ def _video_span(model, x): return None +def _step_position(args): + """(step index, step count) of the block call `args` belongs to. Core puts the whole + schedule in transformer_options["sample_sigmas"] and this step's sigma in ["sigmas"]; + raises when either is missing or the sigma is not on the schedule.""" + to = args.get("transformer_options") or {} + sched, cur = to.get("sample_sigmas"), to.get("sigmas") + if sched is None or cur is None: + raise KeyError("sample_sigmas/sigmas") + cur = cur.reshape(-1)[0].to(sched.dtype) + hit = torch.nonzero(torch.isclose(sched, cur, rtol=1e-4)) + if not len(hit): + raise ValueError("current sigma is not on the schedule") + # sigmas holds N+1 values for N steps (the trailing 0 is never a denoise call). + return int(hit[0]), int(sched.shape[0]) - 1 + + # Set in transformer_options for late-branch guidance's weak pass. Capture hooks skip it: # a deliberately weakened forward is not the network's behaviour to learn from. WEAK_BRANCH_FLAG = "funpack_weak_branch" @@ -3202,6 +3220,18 @@ class FunPackLTXAVSceneChainSampler: "default": 43, "min": 1, "max": 49, "tooltip": "The block the weak copy skips; every block before it is shared. Higher = cheaper and subtler.", }), + "stas": (list(_stas.MODES), { + "default": "off", + "tooltip": "EXPERIMENTAL, H3 only (STAS, arXiv 2603.17825). Over the first 40% of steps, sets the few huge 'massive activation' channels on the first frame and frame edges to alpha x their peak at one block. Aimed at steadier motion; may calm subtle motion. No extra model call. learned = alpha learned from ratings; manual = fixed below.", + }), + "stas_alpha": ("FLOAT", { + "default": 2.0, "min": 0.5, "max": 4.0, "step": 0.1, + "tooltip": "manual mode only. The paper used 1.2-2.5.", + }), + "stas_block": ("INT", { + "default": 15, "min": 0, "max": 49, + "tooltip": "The block whose output is steered. The paper's best was ~30% deep (block 9 of 30).", + }), # A connection socket, never a widget — safe at the end, and it must stay after # every widget above (see the widgets_values note at the top of this block). "second_pass_sigmas": ("SIGMAS", { @@ -4128,6 +4158,13 @@ class FunPackLTXAVSceneChainSampler: # it declines (no negative, both scales 1.0, non-finite knob). _shadow_installed = model is not _model_before_shadow + _stas_installed, _stas_report = False, None + if self._is_h3 and getattr(self, "_stas_alpha", None) is not None: + _model_before_stas = model + model, _stas_report = self._install_stas(model, self._stas_block, + self._stas_alpha, samples) + _stas_installed = model is not _model_before_stas + # Outermost of the block hooks: its weak pass must be able to skip every hook on the # shared blocks, and run the tail exactly as the normal pass does. _late_installed = False @@ -4148,7 +4185,7 @@ class FunPackLTXAVSceneChainSampler: # it silently never ran. # Late-branch guidance's weak pass skips blocks, so its allocations differ too. _compiler_conflict = (_block_repeat_installed or _repr_injecting or _shadow_installed - or _late_installed) + or _late_installed or _stas_installed) try: with self._compiler_disabled_for_forward_patches(_compiler_conflict): sampled = comfy.sample.sample_custom( @@ -4156,6 +4193,8 @@ class FunPackLTXAVSceneChainSampler: noise_mask=latent.get("noise_mask"), seed=int(seed), callback=_progress_cb if pbar is not None else None, ) + if _stas_report is not None: + print(f"[FunPackSceneChain] STAS: {_stas_report()}") if _phrase_probe is not None: self._report_phrase_probe(_phrase_probe[1]) # The masked probe passes ran through every capture hook too, so this @@ -8213,17 +8252,7 @@ class FunPackLTXAVSceneChainSampler: if n == 0: return True try: - to = args.get("transformer_options") or {} - sched, cur = to.get("sample_sigmas"), to.get("sigmas") - if sched is None or cur is None: - raise KeyError("sample_sigmas/sigmas") - cur = cur.reshape(-1)[0].to(sched.dtype) - hit = torch.nonzero(torch.isclose(sched, cur, rtol=1e-4)) - if not len(hit): - raise ValueError("current sigma is not on the schedule") - # sigmas holds N+1 values for N steps (the trailing 0 is never a denoise call). - total = int(sched.shape[0]) - 1 - step_index = int(hit[0]) + step_index, total = _step_position(args) return step_index >= total - n if n > 0 else step_index < -n except Exception as _e: # noqa: BLE001 _log.failed("FunPackSceneChain", "H3 block repeat step window", _e, @@ -8443,6 +8472,96 @@ class FunPackLTXAVSceneChainSampler: _log.failed("FunPackSceneChain", "H3 block repeat", _e, "blocks run once as usual") return model + def _install_stas(self, model, block, alpha, samples): + """STAS (stas.py), H3: steer the massive-activation channels of one block's output, + on the video rows of the first frame and each frame's edges, over the first steps. + -> (model, report fn | None). The model comes back untouched when it cannot install.""" + try: + try: + from . import h3_repr_steering as _rs + except ImportError: + import h3_repr_steering as _rs + net = getattr(getattr(model, "model", None), "diffusion_model", None) + blocks = getattr(net, "blocks", None) + if blocks is None or not 0 <= int(block) < len(blocks): + _log.failed("FunPackSceneChain", "STAS", f"block {block} is not a block of this model", + "the run goes on without it") + return model, None + video = samples.unbind()[0] if getattr(samples, "is_nested", False) else samples + _pt, ph, pw = tuple(getattr(net, "patch_size", (1, 2, 2))) + per_frame = -(-int(video.shape[-2]) // ph) * -(-int(video.shape[-1]) // pw) + patched = model.clone() + to = patched.model_options.get("transformer_options", {}).copy() + patches_replace = dict(to.get("patches_replace", {})) + dit_patches = dict(patches_replace.get("dit", {})) + inner = dit_patches.get(("double_block", int(block))) + _alpha = float(alpha) + cache, stats = {}, {"steered": 0, "empty": 0, "dims": None, "why": None} + + def _stas_hook(args, extra): + out = extra["original_block"](args)["img"] if inner is None else \ + inner(args, extra)["img"] + try: + try: + index, total = _step_position(args) + except Exception as _e: # noqa: BLE001 + stats["why"] = f"no step schedule to read ({_e}), nothing steered" + return {"img": out} + if index >= _stas.first_steps(total): + return {"img": out} + seq_len = int(out.shape[0]) + span = cache.get(seq_len) + if span is None: + mask = _rs.video_mask_from_mod_segments( + args.get("mod_segments"), seq_len, out.device) + idx = mask.nonzero().flatten() if mask is not None else None + span = "none" + if idx is not None and idx.numel(): + a, n = int(idx[0]), int(idx.numel()) + if int(idx[-1]) - a + 1 == n and n % per_frame == 0: + span = (a, n, _stas.target_rows(n // per_frame, per_frame)) + cache[seq_len] = span + if span == "none": + stats["why"] = ("picture rows not found as whole frames in this " + "sequence, nothing steered") + return {"img": out} + a, n, rows = span + video_rows = out[a:a + n] # a view: steered in place + dims, peaks, ratios = _stas.ma_dims(video_rows) + if not dims.numel(): + stats["empty"] += 1 + return {"img": out} + _stas.steer(video_rows, rows, dims, peaks, _alpha) + stats["steered"] += 1 + if stats["dims"] is None: + stats["dims"] = ", ".join(f"{int(d)} ({float(r):.0f}x)" + for d, r in zip(dims, ratios)) + except Exception as _e: # noqa: BLE001 + stats["why"] = f"failed ({_e}), block left as it was" + return {"img": out} + + dit_patches[("double_block", int(block))] = _tag_dit_hook(_stas_hook, inner) + patches_replace["dit"] = dit_patches + to["patches_replace"] = patches_replace + patched.model_options["transformer_options"] = to + + def _report(): + if stats["steered"]: + line = (f"alpha {_alpha:.2f} at block {block}: steered {stats['steered']} " + f"call(s), channels {stats['dims']}") + if stats["empty"]: + line += f"; {stats['empty']} call(s) had no massive channel" + return line + return (stats["why"] or f"no channel at block {block} is over " + f"{_stas.MA_RATIO:.0f}x the mean, nothing steered -- try another block") + + print(f"[FunPackSceneChain] STAS: alpha {_alpha:.2f} at block {block}, first " + f"{_stas.EARLY_FRACTION:.0%} of steps, frame 0 + frame edges, picture only.") + return patched, _report + except Exception as _e: # noqa: BLE001 + _log.failed("FunPackSceneChain", "STAS", _e, "the run goes on without it") + return model, None + def _install_late_branch(self, model, branch, w, n_blocks=50): """Late-branch guidance (late_guidance.py), H3. @@ -8500,9 +8619,14 @@ class FunPackLTXAVSceneChainSampler: old_wrapper = patched.model_options.get("model_function_wrapper") _w = float(w) - def _guided(apply_fn): + def _guided(apply_fn, cond_only): def _apply(x, t, **c): normal = apply_fn(x, t, **c) + if not cond_only: + # A negative-prompt call (CFG above 1): guidance is for the picture + # the prompt asks for, and a weak copy here would double the cost. + saved.pop("h", None) + return normal if "h" not in saved: return normal weak_to = dict(c.get("transformer_options") or {}) @@ -8524,7 +8648,8 @@ class FunPackLTXAVSceneChainSampler: def _late_wrapper(apply_fn, args): # Innermost: every other wrapper sees one guided prediction, not two calls. - guided = _guided(apply_fn) + cond_only = all(int(i) == 0 for i in (args.get("cond_or_uncond") or [0])) + guided = _guided(apply_fn, cond_only) if old_wrapper is not None: return old_wrapper(guided, args) return guided(args["input"], args["timestep"], **args.get("c", {})) @@ -9035,6 +9160,7 @@ class FunPackLTXAVSceneChainSampler: h3_shadow_negative_end_percent=0.60, h3_shadow_negative_compose=False, shot_memory="off", shot_memory_amount=0.7, tsr="off", tsr_k=1.0, late_guidance="off", late_guidance_strength=0.5, late_guidance_block=43, + stas="off", stas_alpha=2.0, stas_block=15, unique_id=None, prompt=None): if not isinstance(positive, list) or not positive: raise ValueError("positive conditioning must contain at least one scene entry.") @@ -9269,12 +9395,12 @@ class FunPackLTXAVSceneChainSampler: # tagged 'funpack_batch_variant'. That marker is the only trigger — the sampler has no # batch-count input. Sample one chain per packed entry, persist each for rating in Studio. if self._split_batch_variants(positive) is not None: - if shot_memory != "off" or tsr != "off" or late_guidance != "off": + if shot_memory != "off" or tsr != "off" or late_guidance != "off" or stas != "off": # Batch training changes one thing at a time; these would change the # noise / the prediction on every variant too. - print("[FunPackSceneChain] shot memory / decisiveness / late-branch guidance " - "are off during batch training -- it varies one setting at a time, and " - "these would vary more") + print("[FunPackSceneChain] shot memory / decisiveness / late-branch guidance / " + "STAS are off during batch training -- it varies one setting at a time, " + "and these would vary more") return self._run_batch_training( model, vae, positive, negative, sampler, sigmas, seed, latent_template, num_frames_per_scene, frame_overlap, cfg, max_scenes, use_same_seed, @@ -9455,10 +9581,13 @@ class FunPackLTXAVSceneChainSampler: self._shot_memory = None self._tsr_k = 1.0 self._late_w, self._late_block = 0.0, int(late_guidance_block) - if shot_memory != "off" or tsr == "learned" or late_guidance == "learned": + self._stas_alpha, self._stas_block = None, int(stas_block) + if (shot_memory != "off" or tsr == "learned" or late_guidance == "learned" + or stas == "learned"): if not refinement_key_input: print("[FunPackSceneChain] shot memory / learned decisiveness / learned " - "late-branch guidance: needs refinement_key_input -- off this run") + "late-branch guidance / learned STAS: needs refinement_key_input -- " + "off this run") if shot_memory != "off" and refinement_key_input: self._shot_memory = _shot_memory.ShotMemory( refinement_key_input, shot_memory, shot_memory_amount) @@ -9470,10 +9599,19 @@ class FunPackLTXAVSceneChainSampler: if not self._is_h3: print("[FunPackSceneChain] late-branch guidance: H3 only -- off this run") else: - self._late_w, _late_note = _late.choose_w( + _w, _late_note = _late.DIAL.choose( refinement_key_input, late_guidance, late_guidance_strength) + self._late_w = max(0.0, _w or 0.0) if _late_note: print(f"[FunPackSceneChain] late-branch guidance: {_late_note}") + if stas != "off": + if not self._is_h3: + print("[FunPackSceneChain] STAS: H3 only -- off this run") + else: + self._stas_alpha, _stas_note = _stas.DIAL.choose( + refinement_key_input, stas, stas_alpha) + if _stas_note: + print(f"[FunPackSceneChain] STAS alpha: {_stas_note}") # Per-bucket value functions for trajectory_guidance. Separate from the one above # because they answer a different question: that one asks "is this a good finish", @@ -10729,9 +10867,13 @@ class FunPackLTXAVSceneChainSampler: else: _tsr.clear_pending(refinement_key_input) if late_guidance == "learned" and self._is_h3: - _late.save_pending(refinement_key_input, self._late_w) + _late.DIAL.save_pending(refinement_key_input, self._late_w) else: - _late.clear_pending(refinement_key_input) + _late.DIAL.clear_pending(refinement_key_input) + if stas == "learned" and self._stas_alpha is not None: + _stas.DIAL.save_pending(refinement_key_input, self._stas_alpha) + else: + _stas.DIAL.clear_pending(refinement_key_input) # Trajectory probe: same run/rating pairing as the two above — the rating that scores # THIS run appends these per-bucket descriptors to the measurement log. diff --git a/stas.py b/stas.py new file mode 100644 index 00000000..bcf0127f --- /dev/null +++ b/stas.py @@ -0,0 +1,66 @@ +"""STAS: steer the massive activations of a video DiT (arXiv 2603.17825). + +A few channels of a video transformer's hidden state carry huge values ("massive +activations"), strongest on the first frame and on the edges of each latent frame. They +act as the model's own rescalers. STAS sets them, at one block and over the first steps +only, to alpha x their current peak: + + g(i, d) = alpha * max_j |H[j, d]| * sign(H[i, d]) for i in S, d in M + +S = every token of the first latent frame + the first and last p% of every frame's +tokens; M = channels whose peak is over 50x the mean magnitude. Paper settings: one +block ~30% deep, first 20 of 50 steps, p = 8%, alpha 1.2-2.5 (tested on Wan, CogVideoX). +No extra model call. Alpha is learned from ratings (rated_dial.py) or set by hand. +""" + +import math + +import torch + +try: + from .rated_dial import Dial, MODES # noqa: F401 +except ImportError: + from rated_dial import Dial, MODES # noqa: F401 + +DIAL = Dial("stas", start=2.0, lo=0.5, hi=3.0, explore=0.2) +DEFAULT_BLOCK = 15 # ~30% of H3's 50, the paper's depth (block 9 of Wan's 30) +EARLY_FRACTION = 0.4 # first 20 of 50 steps +EDGE = 0.08 # p: head and tail share of each frame's tokens +MA_RATIO = 50.0 +MAX_DIMS = 8 + + +def first_steps(total): + """How many of `total` steps are steered.""" + return max(1, math.ceil(EARLY_FRACTION * int(total))) + + +def target_rows(frames, per_frame, edge=EDGE): + """Row indices (within the video rows) of S: all of frame 0, then the head and tail + `edge` share of every frame. Rows are frame-major, as H3 packs them.""" + k = max(1, round(edge * per_frame)) + within = torch.cat([torch.arange(k), torch.arange(per_frame - k, per_frame)]).unique() + rows = (torch.arange(frames)[:, None] * per_frame + within[None, :]).flatten() + return torch.cat([torch.arange(per_frame), rows]).unique() + + +def ma_dims(video): + """-> (channel indices, their peaks, their peak/mean ratios). `video` is [n, D].""" + peak = torch.maximum(video.amax(dim=0).float().abs(), video.amin(dim=0).float().abs()) + stride = max(1, video.shape[0] // 1024) + mean = video[::stride].float().abs().mean().clamp(min=1e-12) + ratio = peak / mean + dims = torch.nonzero(ratio > MA_RATIO).flatten() + if dims.numel() > MAX_DIMS: + dims = dims[ratio[dims].topk(MAX_DIMS).indices] + return dims, peak[dims], ratio[dims] + + +def steer(video, rows, dims, peaks, alpha): + """Set video[rows, dims] to alpha * peak * sign, in place.""" + if not dims.numel() or not rows.numel(): + return video + r, d = rows.to(video.device)[:, None], dims.to(video.device)[None, :] + target = (alpha * peaks.to(video.device))[None, :] * torch.sign(video[r, d].float()) + video[r, d] = target.to(video.dtype) + return video diff --git a/tests/test_late_guidance.py b/tests/test_late_guidance.py index 708e659b..759393ff 100644 --- a/tests/test_late_guidance.py +++ b/tests/test_late_guidance.py @@ -109,19 +109,36 @@ def test_video_span_only_the_picture_is_guided(monkeypatch): assert torch.allclose(out[..., 5:], torch.full((1, 1, 3), 21.0)) # sound: normal pass +def test_negative_prompt_calls_are_not_guided(): + """CFG above 1: the negative call runs once, plain; the prompt call is still guided.""" + import samplers + net = _Blocks() + patched = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(_Model(net), 3, 0.5) + wrapper = patched.model_options["model_function_wrapper"] + x = torch.zeros(1, 2, 1, 2, 2) + c = {"transformer_options": patched.model_options["transformer_options"]} + neg = wrapper(_apply(net), {"input": x, "timestep": torch.tensor([0.5]), "c": c, + "cond_or_uncond": [1]}) + assert net.calls == [0, 1, 2, 3, 4, 5] and torch.allclose(neg, torch.full_like(x, 21.0)) + pos = wrapper(_apply(net), {"input": x, "timestep": torch.tensor([0.5]), "c": c, + "cond_or_uncond": [0]}) + assert torch.allclose(pos, torch.full_like(x, 23.0)) + + def test_strength_learns_from_ratings(monkeypatch, tmp_path): - monkeypatch.setattr(lg, "_path", lambda key: str(tmp_path / f"{key}.late.json")) + d = lg.DIAL + monkeypatch.setattr(d, "path", lambda key: str(tmp_path / f"{key}.late.json")) rng = random.Random(1) for _ in range(30): - w, _note = lg.choose_w("k", "learned", rng=rng) - lg.save_pending("k", w) - lg.commit("k", 1.0 if w > 0.8 else -1.0) - assert lg.learned_w(lg._read("k")["history"]) > 0.8 - assert lg.commit("k", 1.0) is None # nothing pending - lg.save_pending("k", 0.3) - assert lg.commit("k", 0.0) is None # neutral: teaches nothing - lg.save_pending("k", 0.3) - lg.clear_pending("k") # a run with it off - assert lg.commit("k", -1.0) is None - assert lg.choose_w("k", "off")[0] == 0.0 - assert lg.choose_w("", "learned")[0] == 0.0 # no key: off + w, _note = d.choose("k", "learned", 0.5, rng=rng) + d.save_pending("k", w) + d.commit("k", 1.0 if w > 0.8 else -1.0) + assert d.learned(d.read("k")["history"]) > 0.8 + assert d.commit("k", 1.0) is None # nothing pending + d.save_pending("k", 0.3) + assert d.commit("k", 0.0) is None # neutral: teaches nothing + d.save_pending("k", 0.3) + d.clear_pending("k") # a run with it off + assert d.commit("k", -1.0) is None + assert d.choose("k", "off", 0.5)[0] is None + assert d.choose("", "learned", 0.5)[0] is None # no key: off diff --git a/tests/test_stas.py b/tests/test_stas.py new file mode 100644 index 00000000..2ddcdd2e --- /dev/null +++ b/tests/test_stas.py @@ -0,0 +1,97 @@ +"""STAS (stas.py + samplers._install_stas): the right rows and channels are set to +alpha x peak x sign, only on the first steps, only on the picture rows.""" +import sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import stas # noqa: E402 + + +def test_target_rows_are_frame0_plus_each_frames_edges(): + rows = stas.target_rows(frames=3, per_frame=25, edge=0.08) # k = 2 + expect = set(range(25)) | {25, 26, 48, 49, 50, 51, 73, 74} + assert set(rows.tolist()) == expect + + +def test_only_massive_channels_are_found_and_steered(): + v = torch.randn(200, 128) * 0.1 + v[:, 5] = 20.0 * torch.sign(torch.randn(200)) # massive on every token + v[3, 5] = -40.0 # its peak + dims, peaks, _ = stas.ma_dims(v) + assert dims.tolist() == [5] and float(peaks[0]) == 40.0 + before = v.clone() + stas.steer(v, torch.tensor([0, 1, 2]), dims, peaks, 2.0) + assert torch.equal(v[:3, 5], 80.0 * torch.sign(before[:3, 5])) + assert torch.equal(v[3:], before[3:]) and torch.equal(v[:, :5], before[:, :5]) + + +def test_first_steps(): + assert stas.first_steps(50) == 20 and stas.first_steps(7) == 3 and stas.first_steps(1) == 1 + + +class _Net: + patch_size = (1, 2, 2) + + def __init__(self): + self.blocks = [lambda h: h] * 4 + + +class _Model: + def __init__(self): + self.model_options = {"transformer_options": {}} + self.model = type("M", (), {})() + self.model.diffusion_model = _Net() + + def clone(self): + m = _Model() + m.model_options = {k: (dict(v) if isinstance(v, dict) else v) + for k, v in self.model_options.items()} + return m + + +def _call(patched, h, step, total, video_rows): + import samplers + hook = patched.model_options["transformer_options"]["patches_replace"]["dit"][("double_block", 2)] + sched = torch.linspace(1.0, 0.0, total + 1) + args = {"img": h, "mod_segments": video_rows, + "transformer_options": {"sample_sigmas": sched, "sigmas": sched[step:step + 1]}} + return hook(args, {"original_block": lambda a: {"img": a["img"]}})["img"], samplers + + +def test_hook_steers_picture_rows_on_early_steps_only(monkeypatch): + import samplers + import h3_repr_steering as rs + # 10 text rows, then 2 frames x (4x6 latent -> 2x3 = 6 tokens) = 12 video rows, then audio + def mask(video_rows, seq_len, device): + m = torch.zeros(seq_len, dtype=torch.bool) + m[10:22] = True + return m + monkeypatch.setattr(rs, "video_mask_from_mod_segments", mask) + samples = torch.zeros(1, 24, 2, 4, 6) + patched, report = samplers.FunPackLTXAVSceneChainSampler()._install_stas(_Model(), 2, 2.0, samples) + h = torch.randn(30, 64) * 0.1 + h[10:22, 3] = 30.0 + h[12, 3] = 50.0 + out, _ = _call(patched, h.clone(), step=0, total=7, video_rows=None) + # every video row is in S here (frame 0 = all of frame 0; edges k=1 of 6 per frame) + targets = [10, 11, 12, 13, 14, 15, 16, 21] + assert torch.allclose(out[targets, 3], torch.full((8,), 100.0)) + assert torch.equal(out[:10], h[:10]) and torch.equal(out[22:], h[22:]) # text/audio + assert torch.equal(out[17:21], h[17:21]) # interior + late, _ = _call(patched, h.clone(), step=5, total=7, video_rows=None) + assert torch.equal(late, h) # past 40% + assert "channels 3" in report() + + +def test_unreadable_layout_is_declared(monkeypatch): + import samplers + import h3_repr_steering as rs + monkeypatch.setattr(rs, "video_mask_from_mod_segments", lambda *a: None) + patched, report = samplers.FunPackLTXAVSceneChainSampler()._install_stas( + _Model(), 2, 2.0, torch.zeros(1, 24, 2, 4, 6)) + h = torch.randn(30, 8) + out, _ = _call(patched, h.clone(), step=0, total=7, video_rows=None) + assert torch.equal(out, h) and "nothing steered" in report()