diff --git a/samplers.py b/samplers.py index ee58df75..c236dcaa 100644 --- a/samplers.py +++ b/samplers.py @@ -3821,15 +3821,28 @@ class FunPackLTXAVSceneChainSampler: model = self._install_h3_repr_steering( model, refinement_key, h3_repr_steering_strength, _repr_capture, steer_block=h3_repr_steering_block) + model = self._install_h3_av_decouple(model, h3_av_decouple) + model = self._install_h3_attn_temperature( + model, h3_explore_temperature, h3_explore_temperature_block) _q_steer_capture = [{}] if refinement_key and self._parse_block_spec(h3_q_steer_block): - # h3_av_decouple, if also enabled, becomes the OUTER optimized_attention_override - # (installed after this one) and always calls its inner override with a non-None - # mask for exactly the video-row partition -- which is the one thing q_steering's - # own override treats as "someone else already handled this call, back off". The - # two features cannot currently compose: q_steering would capture/inject nothing - # while still paying the clone+hook cost and printing a false "applying" line, so - # skip installing it entirely and say why, rather than let it silently no-op. + # Installed LAST (outermost optimized_attention_override) so it always sees Q + # before av_decouple or attn_temperature touch it -- capture must reflect the + # network's natural Q, not a temperature-scaled one, if h3_q_steer_block and + # h3_explore_temperature_block ever name the same block (each learns/applies + # independently, but they'd otherwise silently compose: temperature's scaling + # would leak into what gets captured and banked as "natural"). + # + # h3_av_decouple still cannot compose with this at all regardless of order: it + # always calls its inner override with a non-None mask for exactly the video-row + # partition, which is the one case q_steering's own override treats as "someone + # else already handled this call, back off" -- so it would capture/inject nothing + # while still paying the clone+hook cost and printing a false "applying" line. + # Skip installing it entirely and say why, rather than let it silently no-op. + # `> 0.0` here, matching _install_h3_av_decouple's own gate (`not (x > 0.0)` to + # bail) on the SAME comparison -- for a NaN strength both must agree on "off", or + # this guard and that gate would independently decide "on"/"off" and reproduce + # the exact bug this comment is describing. if float(h3_av_decouple or 0.0) > 0.0: print("[FunPackSceneChain] H3 query steering: h3_av_decouple is also enabled " "this run -- its attention masking defeats query steering's video-row " @@ -3840,9 +3853,6 @@ class FunPackLTXAVSceneChainSampler: model = self._install_h3_q_steering( model, refinement_key, h3_q_steer_strength, _q_steer_capture, steer_block=h3_q_steer_block) - model = self._install_h3_av_decouple(model, h3_av_decouple) - model = self._install_h3_attn_temperature( - model, h3_explore_temperature, h3_explore_temperature_block) try: sampled = comfy.sample.sample_custom( @@ -6908,7 +6918,12 @@ class FunPackLTXAVSceneChainSampler: _strength = float(strength or 0.0) except (TypeError, ValueError): _strength = 0.0 - if _strength <= 0.0: + # NOT (>0.0), not <=0.0 -- for NaN both comparisons are False under IEEE 754, so + # `<= 0.0` would let a NaN strength through as "on" while _sample_chunk's own + # `float(h3_av_decouple or 0.0) > 0.0` guard (deciding whether to skip installing + # h3_q_steering) treats the SAME NaN as "off" -- the two would disagree and silently + # recreate the av_decouple-defeats-q_steering bug this file already fixed once. + if not (_strength > 0.0): return model try: try: diff --git a/tests/test_h3_av_decouple.py b/tests/test_h3_av_decouple.py index 2a62a222..1f59e2b7 100644 --- a/tests/test_h3_av_decouple.py +++ b/tests/test_h3_av_decouple.py @@ -67,6 +67,17 @@ def test_zero_strength_is_a_true_noop(): assert node._install_h3_av_decouple(model, "not-a-number") is model +def test_nan_strength_is_also_a_true_noop(): + """NaN fails every ordered comparison under IEEE 754, so a bail check written as + `strength <= 0.0` treats NaN as "on" -- the opposite of what `strength > 0.0`-based + checks elsewhere (e.g. the sampler's own guard deciding whether h3_q_steering can be + installed alongside this) would conclude for the same value. Must bail here too, or the + two disagree and silently recreate the av_decouple-defeats-q_steering composition bug.""" + model = _FakeModel() + node = S() + assert node._install_h3_av_decouple(model, float("nan")) is model + + def test_enabled_installs_override_and_block0_hook(): node = S() patched = node._install_h3_av_decouple(_FakeModel(), 0.5) diff --git a/tests/test_h3_q_steering.py b/tests/test_h3_q_steering.py index 81c66350..a278654b 100644 --- a/tests/test_h3_q_steering.py +++ b/tests/test_h3_q_steering.py @@ -298,5 +298,42 @@ def test_q_steer_alone_is_unaffected_by_the_av_decouple_guard(monkeypatch, capsy assert "h3_av_decouple is also enabled" not in out +def test_captures_natural_q_even_when_temperature_targets_the_same_block(monkeypatch): + """h3_q_steer_block and h3_explore_temperature_block naming the SAME block must not leak + temperature's Q-scaling into what q_steering captures/injects -- capture is supposed to + reflect the network's natural Q (same promise h3_repr_steering makes for hidden state). + Requires q_steering to be installed AFTER (more outer than) attn_temperature, matching + the order _sample_chunk now uses -- this test composes them directly in that order.""" + monkeypatch.setattr(_rs, "direction", lambda *a, **k: (None, 0, 0)) + model = S()._install_h3_attn_temperature(_FakeModel(), 1.0, "0") # 2x temperature + capture_holder = [{}] + model = S()._install_h3_q_steering(model, "fake_key", 0.0, capture_holder, steer_block="0") + + # _MOD_SEGMENTS tags BOTH row 0 and row 2 as video -- set both to the SAME value so the + # captured mean-over-video-rows equals that value exactly, cleanly distinguishing + # "natural Q" (5.0) from "already temperature-scaled Q" (2.5) -- a single video row set + # to 5.0 with the other left at 0 would average to 2.5 either way and prove nothing. + q = torch.zeros(1, 2, 4, 1) + q[0, 0, 0, 0] = 5.0 + q[0, 0, 2, 0] = 5.0 + + def fake_func(q_, k_, v_, heads_, mask=None, skip_reshape=True, **kw): + return q_.clone() + + dit = model.model_options["transformer_options"]["patches_replace"]["dit"] + hook = dit[("double_block", 0)] # q_steering's hook, wrapping temperature's + + def _original_block(a): + override = model.model_options["transformer_options"]["optimized_attention_override"] + override(fake_func, q, torch.zeros_like(q), torch.zeros_like(q), 2, skip_reshape=True) + return {"img": a["img"]} + + hook({"img": torch.ones(4, 3), "mod_segments": _MOD_SEGMENTS}, {"original_block": _original_block}) + + # If temperature's 2x scaling had already been applied before q_steering saw Q, this + # would read 2.5, not 5.0. + assert torch.allclose(capture_holder[0][0], torch.tensor([5.0, 0.0])) + + if __name__ == "__main__": print("run via pytest -- these tests need the monkeypatch/capsys fixtures")