From 6c9b0a91aaf759e1e9cc2ef6faf3470cc282a0d5 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:03:56 +0300 Subject: [PATCH] feat: REINS steers multiple blocks at once, block picker covers all 50 h3_repr_steering_block goes from a fixed combo (18 candidate blocks) to free text -- a single block, a range, or a comma list, same syntax as Repeat block(s) (31-40 / 4,5,6). Each named block steers with its OWN learned direction from its OWN liked/disliked history at that block -- there is no shared 'the direction', just N independent ones applied at their own positions in the same forward pass. A block short on rated data at itself just captures, same as any unnamed candidate block. CANDIDATE_BLOCKS widened from 18 curated blocks to all 50 -- every block is now both capturable and steerable, not just the ones on the old grid. Fixed a pre-existing sync gap: h3_block_repeat_last_steps was never added to settings_card.py's ownership/defaults tables, breaking two golden-file tests as soon as they ran (test_settings_card_noise.py). --- h3_repr_steering.py | 14 ++-- movie_editor/backend/settings_card.py | 2 + movie_editor/frontend/engine_settings.js | 9 ++- samplers.py | 78 ++++++++++--------- tests/test_h3_repr_steering.py | 14 ++-- tests/test_h3_repr_steering_multiblock.py | 93 +++++++++++++++++++++++ 6 files changed, 156 insertions(+), 54 deletions(-) create mode 100644 tests/test_h3_repr_steering_multiblock.py diff --git a/h3_repr_steering.py b/h3_repr_steering.py index 0b528b30..b4595854 100644 --- a/h3_repr_steering.py +++ b/h3_repr_steering.py @@ -43,12 +43,14 @@ import torch # of guessed from someone else's percentage. Only DEFAULT_BLOCK ever STEERS; the rest of # CANDIDATE_BLOCKS are captured read-only alongside it so block_sweep() has something to rank. DEFAULT_BLOCK = 25 -# 41-49 individually, not on the coarse grid: the block-influence probe measured the stack as -# heavily BACK-LOADED (blocks 48+49 alone carrying ~a third of all residual movement), so the -# tail is where a steering push has the most to act on and is worth resolving one block at a -# time. Everything below 40 stays coarse -- it is the quiet end. -CANDIDATE_BLOCKS = sorted({1, 5, 10, 15, 20, 25, 30, 35, DEFAULT_BLOCK, - *range(40, 50)}) +# Every block, 0-49: the coarse grid + 41-49 individually was a compromise when only the +# tail's back-loaded weight was known to be worth resolving finely (48+49 alone carrying ~a +# third of all residual movement). 2026-09-04: block 0/1 turned out to be structurally +# distinct too (raw input projection / patch-mixing), not just quiet early blocks -- so +# "coarse in the middle" was an assumption, not a finding. Capturing every block read-only +# costs one extra mean-pool per block per run; steering still only ever touches one block at +# a time (whichever `steer_block` names). +CANDIDATE_BLOCKS = list(range(50)) MIN_PER_GROUP = 2 # need 2+ POSITIVE-weight and 2+ NEGATIVE-weight rows. Was 3 (parity with # absolute/taste steering's floor, not derived from anything specific to this mechanism). diff --git a/movie_editor/backend/settings_card.py b/movie_editor/backend/settings_card.py index f34c163a..8e191e89 100644 --- a/movie_editor/backend/settings_card.py +++ b/movie_editor/backend/settings_card.py @@ -198,6 +198,7 @@ _OWNED_BY = { "dynashift_threshold": ("dynashift", None), "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), "h3_repr_steering_block": ("h3_repr_steering", None), "h3_repr_steering_strength": ("h3_repr_steering", None), "embed_guidance_source": ("embed_guidance", None), @@ -271,6 +272,7 @@ _EDITOR_DEFAULTS = { "h3_block_repeat_times": 1, "h3_block_repeat_span_loop": False, "h3_block_repeat_video_only": False, + "h3_block_repeat_last_steps": 0, "h3_phrase_emphasis": False, "h3_phrase_variability": 0.0, "h3_repr_steering": False, diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index 5579446f..bfdeeac8 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -328,11 +328,12 @@ { name: "dynashift_threshold", label: "DynaShift match threshold", kind: "float", default: 0.6, min: 0.3, max: 0.95, step: 0.05, dependsOn: "dynashift", hint: "How closely a frame must resemble a banked bad one before steering starts. Lower is more aggressive and more likely to push away from content that was actually fine." }, { 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 one block's video-row hidden state each generation, and once you have 3+ liked and 3+ disliked runs on this refinement key, adds their mean difference back into that block's 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." }, + 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", - hint: "Fraction of the block's own hidden-state norm added along the learned direction. Start low and push it deliberately high on a same-seed test to confirm it moves the video at all before trusting a small value." }, - { name: "h3_repr_steering_block", label: "REINS steer at block", kind: "combo", choices: ["1", "5", "10", "15", "20", "25", "30", "35", "40", "41", "42", "43", "44", "45", "46", "47", "48", "49"], default: "25", dependsOn: "h3_repr_steering", - hint: "TEST-ONLY: which of the block-sweep's candidate blocks actually gets steered, instead of always block 25. Blocks 40-49 are listed one at a time because that is where the stack does most of its work (48 and 49 alone carry about a third of all movement). Blocks 1-15 are deep-end territory (the anchor/reference rows haven't been read by the network yet at that depth) -- expect this to behave differently, not just 'weaker', than 20-40. The other candidate blocks are still captured read-only either way." }, + hint: "Fraction of each steered block's own hidden-state norm added along its learned direction — the same strength applies to every block named below. Start low and push it deliberately high on a same-seed test to confirm it moves the video at all before trusting a small value." }, + { name: "h3_repr_steering_block", label: "REINS steer at block(s)", kind: "text", default: "25", dependsOn: "h3_repr_steering", + hint: "TEST-ONLY: which block(s) actually get steered, instead of always block 25 — every block 0-49 is captured read-only regardless. A single block, a range, or a comma list: 25 | 31-40 | 4,5,6 — same syntax as Repeat block(s). Each named block steers with its OWN learned direction from its OWN liked/disliked history at that block; a block without enough data just captures, same as any unnamed candidate block.", + detail: "Blocks 40-49 carry the most residual movement (48+49 alone about a third of it), so a steering push has the most to act on there. Blocks 0-1 are structurally distinct rather than just 'quiet early blocks' — 0 is the raw input projection, 1 does the first cross-patch mixing — expect steering there to behave differently in kind, not just 'weaker', than the middle of the stack." }, { name: "identity_transfer_enabled", label: "Best-FaceID compatibility", kind: "bool", default: false, hint: "Feeds the identity pin image the way Best-FaceID identity LoRAs expect it. Needs an Identity pin set.", detail: "Replaces Continuity's identity-pin guide with separate, non-rendered reference tokens plus an optional ArcFace projector below. Load the LoRA itself in Models. No effect without an identity pin." }, diff --git a/samplers.py b/samplers.py index 0c595e96..b6202686 100644 --- a/samplers.py +++ b/samplers.py @@ -2966,9 +2966,9 @@ class FunPackLTXAVSceneChainSampler: # workflow (builder.py), silently corrupting anything the user hadn't # explicitly overridden in sampler_inputs. Fixed 2026-09-03; do not move this # back up next to its sibling widgets. - "h3_repr_steering_block": (["1", "5", "10", "15", "20", "25", "30", "35", "40", "41", "42", "43", "44", "45", "46", "47", "48", "49"], { + "h3_repr_steering_block": ("STRING", { "default": "25", - "tooltip": "TEST-ONLY: which block actually gets steered (the rest of the sweep's candidate blocks are still captured read-only, same as always). direction() is recomputed for whichever block you pick here -- this does not change what block_sweep ranks, only which block's learned direction gets injected into this run.", + "tooltip": "TEST-ONLY: which block(s) actually get steered (the rest of the sweep's candidate blocks are still captured read-only, same as always). A single block, a range, or a comma list -- 25 | 31-40 | 4,5,6 -- same syntax as h3_block_repeat. Each named block steers with its OWN learned direction, computed independently from its OWN liked/disliked history at that block; this does not change what block_sweep ranks, only which block(s) get injected into this run.", }), "h3_block_repeat": ("STRING", { "default": "", @@ -3790,7 +3790,7 @@ class FunPackLTXAVSceneChainSampler: if h3_repr_steering and refinement_key: model = self._install_h3_repr_steering( model, refinement_key, h3_repr_steering_strength, _repr_capture, - steer_block=int(h3_repr_steering_block)) + steer_block=h3_repr_steering_block) try: sampled = comfy.sample.sample_custom( @@ -6714,33 +6714,41 @@ class FunPackLTXAVSceneChainSampler: steer_block=None): """EXPERIMENTAL, unvalidated (H3 only). See h3_repr_steering.py. - Reaches inside the forward pass: hooks `steer_block` (default DEFAULT_BLOCK), adds the - learned liked-minus-disliked direction to that block's VIDEO rows (never audio or - text — derived fresh from `mod_segments` each call, cached by seq_len). Every OTHER - block in CANDIDATE_BLOCKS is hooked read-only (captured, never modified) at the same - time, so `h3_repr_steering.block_sweep()` has real per-block data to rank without - needing its own separate collection pass. `capture_holder[0]` becomes {block: - descriptor} for the caller to save as the next pending row once sampling finishes. - No-ops on STEERING (clears nothing, costs nothing extra) when there are not yet 2+ - liked and 2+ disliked runs on this key at `steer_block` — direction() says so, this - just reports it; capture at every candidate block still happens regardless. + Reaches inside the forward pass: hooks every block named in `steer_block` (default + DEFAULT_BLOCK), adds EACH block's OWN learned liked-minus-disliked direction to that + block's VIDEO rows (never audio or text — derived fresh from `mod_segments` each call, + cached by seq_len). Every OTHER block in CANDIDATE_BLOCKS is hooked read-only + (captured, never modified) at the same time, so `h3_repr_steering.block_sweep()` has + real per-block data to rank without needing its own separate collection pass. + `capture_holder[0]` becomes {block: descriptor} for the caller to save as the next + pending row once sampling finishes. A block with fewer than 2+ liked and 2+ disliked + runs of its OWN just captures, same as any other candidate block — no-ops on + STEERING there (clears nothing, costs nothing extra), reported once per block. - `steer_block` lets test runs pick which candidate block actually gets injected instead - of always DEFAULT_BLOCK -- direction() is computed fresh for that block. Does not + `steer_block` accepts a single block ("25"), a range ("31-40"), or a comma list + ("4,5,6") — same syntax as h3_block_repeat, parsed by `_parse_block_spec`. Each named + block steers with its OWN direction(), computed independently: two blocks' liked/ + disliked history are unrelated, so there is no shared "the direction" here, only N + separate ones applied at their own positions in the same forward pass. Does not change what block_sweep ranks (still every candidate block, read-only capture as - always); only changes which one this particular run steers with. + always); only changes which ones this particular run steers with. """ try: try: from . import h3_repr_steering as _rs except ImportError: import h3_repr_steering as _rs - steer_block = _rs.DEFAULT_BLOCK if steer_block is None else int(steer_block) - direction, n_liked, n_disliked = _rs.direction(refinement_key, block=steer_block) - if direction is None: - print(f"[FunPackSceneChain] H3 representation steering: not enough data yet " - f"({n_liked} liked / {n_disliked} disliked, need {_rs.MIN_PER_GROUP} of " - f"each at block {steer_block}) — capturing this run, not steering it.") + steer_blocks = (self._parse_block_spec(steer_block) + if steer_block else {_rs.DEFAULT_BLOCK}) + directions = {} + for _sb in sorted(steer_blocks): + direction, n_liked, n_disliked = _rs.direction(refinement_key, block=_sb) + directions[_sb] = direction + if direction is None: + print(f"[FunPackSceneChain] H3 representation steering: not enough data " + f"yet ({n_liked} liked / {n_disliked} disliked, need " + f"{_rs.MIN_PER_GROUP} of each at block {_sb}) — capturing this run, " + "not steering it.") patched = model.clone() to = patched.model_options.get("transformer_options", {}).copy() @@ -6748,7 +6756,7 @@ class FunPackLTXAVSceneChainSampler: dit_patches = dict(patches_replace.get("dit", {})) _strength = float(strength) - def _make_hook(block, inject): + def _make_hook(block, direction): inner = dit_patches.get(("double_block", block)) mask_cache = {} @@ -6762,18 +6770,15 @@ class FunPackLTXAVSceneChainSampler: args.get("mod_segments"), seq_len, out.device) mask_cache[seq_len] = mask if mask is not None: - # Captured BEFORE injection, deliberately -- this is the ONLY block - # where `inject` can be True, so it's the only one where capturing - # the post-injection state would make the logged descriptor depend - # on that run's own strength setting instead of reflecting the - # network's natural behavior. Every other candidate block never - # injects, so this line changes nothing for them; here it's what - # keeps block 25's block_sweep entry comparable across different - # strengths, and comparable to the other 8 blocks in the first place. + # Captured BEFORE injection, deliberately -- a steered block is the + # only kind where capturing the post-injection state would make the + # logged descriptor depend on that run's own strength setting instead + # of reflecting the network's natural behavior. A block that never + # injects (direction is None here) is unaffected by this ordering. desc = _rs.capture(out, mask) if desc is not None: capture_holder[0][block] = desc - if inject and direction is not None and _strength > 0.0: + if direction is not None and _strength > 0.0: rows = out[mask] row_norm = rows.detach().float().norm(dim=-1).mean() if torch.isfinite(row_norm) and row_norm > 0: @@ -6785,14 +6790,15 @@ class FunPackLTXAVSceneChainSampler: for _block in _rs.CANDIDATE_BLOCKS: dit_patches[("double_block", _block)] = _make_hook( - _block, inject=(_block == steer_block)) + _block, directions.get(_block)) patches_replace["dit"] = dit_patches to["patches_replace"] = patches_replace patched.model_options["transformer_options"] = to - if direction is not None: + _steered = sorted(b for b, d in directions.items() if d is not None) + if _steered: print(f"[FunPackSceneChain] H3 representation steering: applying learned " - f"direction ({n_liked} liked / {n_disliked} disliked) at block " - f"{steer_block}, strength {_strength:g}.") + f"direction(s) at block(s) {', '.join(str(b) for b in _steered)}, " + f"strength {_strength:g}.") return patched except Exception as _e: # noqa: BLE001 _log.failed("FunPackSceneChain", "H3 representation steering", _e, diff --git a/tests/test_h3_repr_steering.py b/tests/test_h3_repr_steering.py index d0b4c2ae..1bcc78e9 100644 --- a/tests/test_h3_repr_steering.py +++ b/tests/test_h3_repr_steering.py @@ -227,17 +227,15 @@ def test_block_sweep_drops_neutral_weight_rows(): def test_every_steerable_block_is_also_captured(): """A block can only be STEERED if it is also in CANDIDATE_BLOCKS -- direction() reads - descriptors captured at that exact block, so a picker entry outside the capture set would - report "not enough data yet" forever no matter how many runs were rated.""" - import re - src = open("samplers.py").read() - choices = re.search(r'"h3_repr_steering_block": \(\[([^]]*)\]', src).group(1) - offered = {int(c.strip().strip('"')) for c in choices.split(",")} - assert offered <= set(rs.CANDIDATE_BLOCKS), offered - set(rs.CANDIDATE_BLOCKS) + descriptors captured at that exact block, so a text-field entry outside the capture set + would report "not enough data yet" forever no matter how many runs were rated. The steer + field is free text now (any block 0-49, or a range/comma list of them via + _parse_block_spec), so this only holds if CANDIDATE_BLOCKS covers the WHOLE valid range.""" + assert set(range(50)) <= set(rs.CANDIDATE_BLOCKS) def test_the_back_loaded_tail_is_resolved_one_block_at_a_time(): """The influence probe measured blocks 48-49 carrying ~a third of all residual movement, so the tail is where a push has the most to act on -- it is covered individually rather - than on the coarse grid the quiet early half uses.""" + than on the coarse grid the quiet early half used to use.""" assert set(range(40, 50)) <= set(rs.CANDIDATE_BLOCKS) diff --git a/tests/test_h3_repr_steering_multiblock.py b/tests/test_h3_repr_steering_multiblock.py new file mode 100644 index 00000000..7c93b060 --- /dev/null +++ b/tests/test_h3_repr_steering_multiblock.py @@ -0,0 +1,93 @@ +"""REINS steer_block went from a single combo choice to free text ("31-40", "4,5,6"), +each named block steering with its OWN learned direction. Exercised at the install level: +does the resulting hook actually inject at every named block and NOWHERE else.""" +import sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +sys.path.insert(0, ".") +import _comfy_stubs # noqa: E402, F401 +from samplers import FunPackLTXAVSceneChainSampler as S # noqa: E402 +import h3_repr_steering as rs # noqa: E402 + + +class _FakeModel: + def __init__(self): + self.model_options = {} + + def clone(self): + m = _FakeModel() + m.model_options = dict(self.model_options) + return m + + +def _run_block(patched, block, seq_len=4): + """Invoke the installed hook for `block` and return the resulting "img" tensor.""" + dit = patched.model_options["transformer_options"]["patches_replace"]["dit"] + hook = dit[("double_block", block)] + img = torch.ones(seq_len, 3) + args = {"img": img, "mod_segments": [(0, seq_len, 0)]} # all rows tagged video + extra = {"original_block": lambda a: {"img": a["img"]}} + return hook(args, extra)["img"] + + +def test_multiple_named_blocks_each_steer_with_their_own_direction(monkeypatch): + directions = {5: torch.tensor([1.0, 0.0, 0.0]), 10: torch.tensor([0.0, 1.0, 0.0])} + + def fake_direction(_key, block=None): + d = directions.get(block) + return (d, 5, 5) if d is not None else (None, 0, 0) + + monkeypatch.setattr(rs, "direction", fake_direction) + node = S() + patched = node._install_h3_repr_steering( + _FakeModel(), "key", strength=1.0, capture_holder=[{}], steer_block="5,10") + + out5 = _run_block(patched, 5) + out10 = _run_block(patched, 10) + out6 = _run_block(patched, 6) # captured, never steered + + assert not torch.allclose(out5, torch.ones(4, 3)), "block 5 should have been steered" + assert not torch.allclose(out10, torch.ones(4, 3)), "block 10 should have been steered" + assert torch.allclose(out6, torch.ones(4, 3)), "an unnamed block must stay untouched" + # Different learned directions -> different injected deltas. + assert not torch.allclose(out5, out10) + + +def test_a_range_steers_every_block_in_it(monkeypatch): + monkeypatch.setattr(rs, "direction", lambda _k, block=None: (torch.tensor([1.0, 0.0, 0.0]), 5, 5)) + node = S() + patched = node._install_h3_repr_steering( + _FakeModel(), "key", strength=1.0, capture_holder=[{}], steer_block="7-9") + + for b in (7, 8, 9): + assert not torch.allclose(_run_block(patched, b), torch.ones(4, 3)), b + assert torch.allclose(_run_block(patched, 6), torch.ones(4, 3)) + assert torch.allclose(_run_block(patched, 10), torch.ones(4, 3)) + + +def test_a_block_without_enough_rated_data_only_captures(monkeypatch): + monkeypatch.setattr(rs, "direction", lambda _k, block=None: (None, 1, 0)) + node = S() + capture_holder = [{}] + patched = node._install_h3_repr_steering( + _FakeModel(), "key", strength=1.0, capture_holder=capture_holder, steer_block="5") + + out = _run_block(patched, 5) + assert torch.allclose(out, torch.ones(4, 3)), "no direction yet -- must not inject" + assert 5 in capture_holder[0], "capture should still happen regardless of steering" + + +def test_empty_steer_block_falls_back_to_default_block(): + node = S() + patched = node._install_h3_repr_steering( + _FakeModel(), "", strength=1.0, capture_holder=[{}], steer_block="") + dit = patched.model_options["transformer_options"]["patches_replace"]["dit"] + assert ("double_block", rs.DEFAULT_BLOCK) in dit + + +if __name__ == "__main__": + test_empty_steer_block_falls_back_to_default_block() + print("ok (run via pytest for the monkeypatch-dependent cases)")