From 608a97a48de4732ff85c6b45966e2956aaf12a00 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Mon, 7 Sep 2026 03:55:56 +0300 Subject: [PATCH] fix: H3 query steering was on by default, and could crash or misreport Three real faults found by adversarial review of 299941a: - sample()'s own default for h3_q_steer_block was "0,1" while INPUT_TYPES declared "" -- any saved workflow predating this widget omits the key entirely, so ComfyUI falls back to sample()'s default, silently turning on capture/clone/hook for every H3 generation with a refinement key. Now both default to blank, matching every sibling knob in this file. - The "applying learned direction(s)" console line fired whenever a direction was available, not when it was actually applied -- strength=0 still printed it, which is the exact silent-success-message trap this project has already been burned by. Now split into an "applying" line (strength>0) and a "capturing only" line (strength=0). - A direction banked under a different heads*head_dim layout (model swap on the same refinement_key) hit an unguarded .view() and would hard-crash generation instead of degrading. Now checked and skipped with a one-time notice, same pattern as the existing batch>1 guard. Added regression tests for all three, plus two coverage-gap tests the review flagged (kind isolation against REAL persistence, not mocks; two named blocks steering with independent directions in one run). --- samplers.py | 34 ++++++++-- tests/test_h3_q_steering.py | 125 ++++++++++++++++++++++++++++++++++++ 2 files changed, 155 insertions(+), 4 deletions(-) diff --git a/samplers.py b/samplers.py index bbaba728..7a164c36 100644 --- a/samplers.py +++ b/samplers.py @@ -7152,6 +7152,7 @@ class FunPackLTXAVSceneChainSampler: _seg_holder = {"mod_segments": None} _active_block = {"cur": None} _warned_batch = {"on": False} + _warned_shape = set() def _make_hook(block): inner = dit_patches.get(("double_block", block)) @@ -7214,6 +7215,21 @@ class FunPackLTXAVSceneChainSampler: direction = directions.get(block) if direction is not None and _strength > 0.0: head_dim = q.shape[-1] + if direction.numel() != heads * head_dim: + # A direction banked under a different head/head_dim layout (model + # swap, a different H3 variant on the same refinement_key) cannot be + # reshaped onto this run's Q at all -- .view() would hard-crash the + # generation instead of degrading, exactly what this file's other + # experimental knobs (batch>1 above, av_decouple's shape check) avoid. + if block not in _warned_shape: + _warned_shape.add(block) + print(f"[FunPackSceneChain] H3 query steering: block {block}'s " + f"learned direction has {direction.numel()} values, but this " + f"run's Q is {heads}x{head_dim} ({heads * head_dim}) -- the " + f"model or its head layout changed since this direction was " + f"captured. Not applying it this run (still capturing).") + return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape, + **kwargs) rows = q[:, :, vmask, :] row_norm = rows.detach().float().norm(dim=-1).mean() if torch.isfinite(row_norm) and row_norm > 0: @@ -7225,12 +7241,22 @@ class FunPackLTXAVSceneChainSampler: to["optimized_attention_override"] = _override patched.model_options["transformer_options"] = to - _steered = sorted(b for b, d in directions.items() if d is not None) - if _steered: + _have_direction = sorted(b for b, d in directions.items() if d is not None) + # Gated on strength too, not just direction-availability -- a learned direction + # sitting at strength 0 injects nothing (see the strength>0.0 check in _override), + # so claiming "applying" there would be exactly the silent-success-message trap + # this project has been burned by before. + if _have_direction and _strength > 0.0: _per_block = ", ".join(f"{b} ({_counts[b][0]}/{_counts[b][1]})" - for b in _steered) + for b in _have_direction) print(f"[FunPackSceneChain] H3 query steering: applying learned direction(s) " f"at block (liked/disliked): {_per_block}, strength {_strength:g}.") + elif _have_direction: + _per_block = ", ".join(f"{b} ({_counts[b][0]}/{_counts[b][1]})" + for b in _have_direction) + print(f"[FunPackSceneChain] H3 query steering: learned direction(s) available " + f"at block (liked/disliked): {_per_block}, but strength is 0 -- " + f"capturing only, nothing applied.") return patched except Exception as _e: # noqa: BLE001 _log.failed("FunPackSceneChain", "H3 query steering", _e, @@ -7887,7 +7913,7 @@ class FunPackLTXAVSceneChainSampler: h3_block_repeat="", h3_block_repeat_times=1, h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False, h3_block_repeat_last_steps=0, - h3_q_steer_strength=0.0, h3_q_steer_block="0,1", + h3_q_steer_strength=0.0, h3_q_steer_block="", alg_guide_blur_strength=2.0, alg_guide_blur_sigma_threshold=0.975, alg_anchor=False, alg_anchor_strength=2.0, alg_anchor_sigma_threshold=0.975, identity_transfer_enabled=False, identity_projector="None", source_id=2.0, diff --git a/tests/test_h3_q_steering.py b/tests/test_h3_q_steering.py index bb67ed92..b7d224e0 100644 --- a/tests/test_h3_q_steering.py +++ b/tests/test_h3_q_steering.py @@ -137,5 +137,130 @@ def test_batch_greater_than_one_is_a_noop(monkeypatch, capsys): assert "batch size > 1" in capsys.readouterr().out +def test_kind_isolation_uses_real_persistence(monkeypatch, tmp_path): + """Real save_pending/commit/direction (no mocking of the persistence layer itself) for + BOTH kind="repr_steer" (REINS' own default) and kind="q_steer" at the SAME block index, + same refinement_key -- must never collide, even though a block index means a different + vector space in each store.""" + monkeypatch.setattr(_rs, "state_path", + lambda key, kind="repr_steer": str(tmp_path / f"{key}.{kind}.pt")) + key = "fake_key" + block = 0 + # REINS' own (hidden-state) direction: liked cluster near [1,0], disliked near [-1,0]. + for i in range(3): + _rs.save_pending(key, {block: torch.tensor([1.0 + i * 0.01, 0.0])}) + _rs.commit(key, reward=1.0) + _rs.save_pending(key, {block: torch.tensor([-1.0 - i * 0.01, 0.0])}) + _rs.commit(key, reward=-1.0) + # q_steer's own direction at the SAME block index: liked/disliked swapped on axis 2 + # instead, so if the two stores ever shared a file/rows the directions would collide. + for i in range(3): + _rs.save_pending(key, {block: torch.tensor([0.0, 1.0 + i * 0.01])}, kind="q_steer") + _rs.commit(key, reward=1.0, kind="q_steer") + _rs.save_pending(key, {block: torch.tensor([0.0, -1.0 - i * 0.01])}, kind="q_steer") + _rs.commit(key, reward=-1.0, kind="q_steer") + + repr_dir, r_pos, r_neg = _rs.direction(key, block=block) + q_dir, q_pos, q_neg = _rs.direction(key, block=block, kind="q_steer") + assert r_pos == 3 and r_neg == 3 + assert q_pos == 3 and q_neg == 3 + # REINS' direction points along axis 1, q_steer's along axis 2 -- if the stores collided + # (same file, or rows merged) both would reflect a mix of both axes instead. + assert repr_dir[0].abs() > 0.9 and repr_dir[1].abs() < 0.1 + assert q_dir[1].abs() > 0.9 and q_dir[0].abs() < 0.1 + + +def test_two_named_blocks_steer_with_independent_directions(monkeypatch): + """steer_block='0,1' -- each block must apply its OWN direction, not the other's.""" + directions = {0: torch.tensor([1.0, 0.0]), 1: torch.tensor([0.0, 1.0])} + + def _fake_direction(refinement_key, block=None, kind="repr_steer"): + assert kind == "q_steer" + return directions[block], 5, 5 + monkeypatch.setattr(_rs, "direction", _fake_direction) + capture_holder = [{}] + patched = S()._install_h3_q_steering( + _FakeModel(), "fake_key", 1.0, capture_holder, steer_block="0,1") + + def fake_func(q_, k_, v_, heads_, mask=None, skip_reshape=True, **kw): + return q_.clone() + + results = {} + for block in (0, 1): + q = torch.zeros(1, 2, 4, 1) + q[0, 0, 0, 0] = 5.0 # heads=2, head_dim=1; masked (video) row 0 only for simplicity + results[block] = _run_attention_inside_block( + patched, block=block, + attn_fn=lambda q=q: _override_fn(patched)( + fake_func, q, torch.zeros_like(q), torch.zeros_like(q), 2, skip_reshape=True)) + + # Block 0's direction is [1,0] -> only head 0 shifts; block 1's is [0,1] -> only head 1 + # shifts. If the blocks shared one direction (a wiring bug), both would shift identically. + assert results[0][0, 0, 0, 0] != 5.0 + assert torch.allclose(results[0][0, 1, 0, 0], torch.tensor(0.0)) + assert torch.allclose(results[1][0, 0, 0, 0], torch.tensor(5.0)) + assert results[1][0, 1, 0, 0] != 0.0 + assert 0 in capture_holder[0] and 1 in capture_holder[0] + + +def test_sample_signature_default_matches_input_types_default(): + """A ComfyUI 'optional' widget that predates a saved workflow is simply ABSENT from the + call -- sample()'s own Python default fires instead of whatever INPUT_TYPES declares. + h3_q_steer_block must default to blank (off) in BOTH places, or every pre-existing + workflow silently turns this on (caught live: sample()'s default was "0,1" while + INPUT_TYPES said "" -- an editor-side workflow, or ANY graph saved before this widget + existed, would have started capturing/hooking/cloning on every H3 generation).""" + import inspect + sig_default = inspect.signature(S.sample).parameters["h3_q_steer_block"].default + input_types_default = S.INPUT_TYPES()["optional"]["h3_q_steer_block"][1]["default"] + assert sig_default == input_types_default == "", ( + f"sample() default {sig_default!r} != INPUT_TYPES default {input_types_default!r} " + "-- must both be blank") + strength_sig_default = inspect.signature(S.sample).parameters["h3_q_steer_strength"].default + strength_it_default = S.INPUT_TYPES()["optional"]["h3_q_steer_strength"][1]["default"] + assert strength_sig_default == strength_it_default == 0.0 + + +def test_strength_zero_does_not_claim_to_be_applying(monkeypatch, capsys): + """A learned direction existing (enough rating history) is not the same as it being + APPLIED -- injection is separately gated on strength > 0. The console message must not + claim "applying" when strength is 0, or a passively-accumulated rating history would + print a false positive on every subsequent generation.""" + patched, capture_holder = _install(monkeypatch, torch.tensor([1.0, -1.0]), strength=0.0) + q = torch.ones(1, 2, 4, 1) + + def fake_func(q_, k_, v_, heads_, mask=None, skip_reshape=True, **kw): + return q_.clone() + + out = _run_attention_inside_block( + patched, block=0, + attn_fn=lambda: _override_fn(patched)( + fake_func, q, torch.zeros_like(q), torch.zeros_like(q), 2, skip_reshape=True)) + assert torch.allclose(out, q) # strength 0 -> no injection despite a real direction + stdout = capsys.readouterr().out + assert "applying learned direction" not in stdout + assert "strength is 0" in stdout + + +def test_mismatched_direction_length_does_not_crash(monkeypatch, capsys): + """A direction banked under a different heads*head_dim layout (a model swap on the same + refinement_key, or a future H3 variant) cannot be reshaped onto this run's Q -- must + degrade with a one-time notice, not raise, the same way batch>1 does above.""" + # direction has 3 values; this run's Q is heads=2, head_dim=1 -> needs exactly 2. + patched, capture_holder = _install(monkeypatch, torch.tensor([1.0, 2.0, 3.0]), strength=1.0) + q = torch.ones(1, 2, 4, 1) * 5.0 + + def fake_func(q_, k_, v_, heads_, mask=None, skip_reshape=True, **kw): + return q_.clone() + + out = _run_attention_inside_block( + patched, block=0, + attn_fn=lambda: _override_fn(patched)( + fake_func, q, torch.zeros_like(q), torch.zeros_like(q), 2, skip_reshape=True)) + assert torch.allclose(out, q) # not applied, but did not crash + assert 0 in capture_holder[0] # capture still happens + assert "learned direction has 3 values" in capsys.readouterr().out + + if __name__ == "__main__": print("run via pytest -- these tests need the monkeypatch/capsys fixtures")