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).
This commit is contained in:
+30
-4
@@ -7152,6 +7152,7 @@ class FunPackLTXAVSceneChainSampler:
|
|||||||
_seg_holder = {"mod_segments": None}
|
_seg_holder = {"mod_segments": None}
|
||||||
_active_block = {"cur": None}
|
_active_block = {"cur": None}
|
||||||
_warned_batch = {"on": False}
|
_warned_batch = {"on": False}
|
||||||
|
_warned_shape = set()
|
||||||
|
|
||||||
def _make_hook(block):
|
def _make_hook(block):
|
||||||
inner = dit_patches.get(("double_block", block))
|
inner = dit_patches.get(("double_block", block))
|
||||||
@@ -7214,6 +7215,21 @@ class FunPackLTXAVSceneChainSampler:
|
|||||||
direction = directions.get(block)
|
direction = directions.get(block)
|
||||||
if direction is not None and _strength > 0.0:
|
if direction is not None and _strength > 0.0:
|
||||||
head_dim = q.shape[-1]
|
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, :]
|
rows = q[:, :, vmask, :]
|
||||||
row_norm = rows.detach().float().norm(dim=-1).mean()
|
row_norm = rows.detach().float().norm(dim=-1).mean()
|
||||||
if torch.isfinite(row_norm) and row_norm > 0:
|
if torch.isfinite(row_norm) and row_norm > 0:
|
||||||
@@ -7225,12 +7241,22 @@ class FunPackLTXAVSceneChainSampler:
|
|||||||
|
|
||||||
to["optimized_attention_override"] = _override
|
to["optimized_attention_override"] = _override
|
||||||
patched.model_options["transformer_options"] = to
|
patched.model_options["transformer_options"] = to
|
||||||
_steered = sorted(b for b, d in directions.items() if d is not None)
|
_have_direction = sorted(b for b, d in directions.items() if d is not None)
|
||||||
if _steered:
|
# 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]})"
|
_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) "
|
print(f"[FunPackSceneChain] H3 query steering: applying learned direction(s) "
|
||||||
f"at block (liked/disliked): {_per_block}, strength {_strength:g}.")
|
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
|
return patched
|
||||||
except Exception as _e: # noqa: BLE001
|
except Exception as _e: # noqa: BLE001
|
||||||
_log.failed("FunPackSceneChain", "H3 query steering", _e,
|
_log.failed("FunPackSceneChain", "H3 query steering", _e,
|
||||||
@@ -7887,7 +7913,7 @@ class FunPackLTXAVSceneChainSampler:
|
|||||||
h3_block_repeat="", h3_block_repeat_times=1,
|
h3_block_repeat="", h3_block_repeat_times=1,
|
||||||
h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False,
|
h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False,
|
||||||
h3_block_repeat_last_steps=0,
|
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_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,
|
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,
|
identity_transfer_enabled=False, identity_projector="None", source_id=2.0,
|
||||||
|
|||||||
@@ -137,5 +137,130 @@ def test_batch_greater_than_one_is_a_noop(monkeypatch, capsys):
|
|||||||
assert "batch size > 1" in capsys.readouterr().out
|
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__":
|
if __name__ == "__main__":
|
||||||
print("run via pytest -- these tests need the monkeypatch/capsys fixtures")
|
print("run via pytest -- these tests need the monkeypatch/capsys fixtures")
|
||||||
|
|||||||
Reference in New Issue
Block a user