Files
olivv-cs-ComfyUI-FunPack/tests/test_block_influence.py
T
digital-garbage be5579f1f9 fix: rank blocks by share of movement, not by a ratio depth decides
Live data showed block 1 at 6.01 and block 0 at 1.83 against ~0.2-0.35
everywhere else, with block 1's liked/disliked gap (+5.82) as large as
its whole measurement. That is not "the first two blocks do 30x the
work" — ||f_i||/||x_i|| divides by a residual stream that GROWS with
depth, so early blocks score high because the denominator is barely
formed. Ranking on it just sorts the earliest blocks to the top, and a
high-variance early block fakes a large taste correlation at low n.

The hook now stores both norms instead of their ratio, and profile()
adds `share` — each block's raw ||f_i|| over every block's, one
denominator for the whole stack. The panel ranks and displays that
("2.3% of the movement"). The ratio is kept, since it is still the right
answer to "how much did this block move the stream relative to what the
stream was".

Checked against data shaped like the report: ratio ranks 0,1,2 (all
artifact), share ranks the blocks that actually move more.
2026-09-04 18:10:55 +03:00

545 lines
20 KiB
Python

"""Per-block influence probe: the pending->rating pairing, the profile maths, and the
flatness readout that decides whether rating-driven block weighting is worth pursuing."""
import sys
import pytest
import torch
sys.path.insert(0, ".")
import block_influence as bi # noqa: E402
@pytest.fixture(autouse=True)
def _isolate_state(monkeypatch, tmp_path):
monkeypatch.setattr(bi, "state_path", lambda key: str(tmp_path / f"{key}.pt"))
def _run(key, profile, reward):
bi.save_pending(key, profile)
return bi.commit(key, reward)
# --- pending / commit contract -------------------------------------------------
def test_commit_reports_recorded_on_a_real_pairing():
assert _run("k", {0: 0.1, 1: 0.2}, 1.0) == "recorded"
def test_commit_reports_no_pending_with_nothing_to_pair():
assert bi.commit("k", 1.0) == "no_pending"
def test_commit_reports_no_key():
assert bi.commit("", 1.0) == "no_key"
def test_pending_overwrites_rather_than_accumulating():
bi.save_pending("k", {0: 0.1})
bi.save_pending("k", {0: 0.9}) # a second run before any rating
bi.commit("k", 1.0)
rows = bi._load("k")["rows"]
assert len(rows) == 1 and rows[0]["profile"][0] == 0.9
def test_clear_all_removes_the_state():
_run("k", {0: 0.1}, 1.0)
bi.clear_all("k")
assert bi.profile("k")["overall"] == {}
def test_no_key_is_a_safe_no_op():
bi.save_pending("", {0: 0.1})
bi.commit("", 1.0)
bi.clear_all("")
# --- profile maths -------------------------------------------------------------
def test_overall_is_the_mean_across_runs_regardless_of_rating():
_run("k", {0: 0.2, 1: 0.4}, 1.0)
_run("k", {0: 0.4, 1: 0.6}, -1.0)
got = bi.profile("k")["overall"]
assert got[0] == pytest.approx(0.3)
assert got[1] == pytest.approx(0.5)
def test_difference_needs_min_per_group_on_each_side():
for _ in range(bi.MIN_PER_GROUP):
_run("k", {0: 0.5}, 1.0)
_run("k", {0: 0.1}, -1.0) # only one disliked
p = bi.profile("k")
assert p["difference"] is None
assert p["n_liked"] == bi.MIN_PER_GROUP and p["n_disliked"] == 1
def test_difference_is_liked_mean_minus_disliked_mean():
for _ in range(bi.MIN_PER_GROUP):
_run("k", {0: 0.8, 1: 0.2}, 1.0)
_run("k", {0: 0.4, 1: 0.2}, -1.0)
diff = bi.profile("k")["difference"]
assert diff[0] == pytest.approx(0.4) # block 0 runs hotter on liked runs
assert diff[1] == pytest.approx(0.0) # block 1 is indifferent to the rating
def test_a_block_missing_from_a_row_is_skipped_not_counted_as_zero():
"""A run that hooked fewer blocks must not drag the absent block's mean toward 0 --
the same rule h3_repr_steering.direction() follows for a missing descriptor."""
_run("k", {0: 0.4, 1: 0.4}, 1.0)
_run("k", {0: 0.6}, 1.0) # no block 1 this run
got = bi.profile("k")["overall"]
assert got[0] == pytest.approx(0.5)
assert got[1] == pytest.approx(0.4)
def test_neutral_weight_rows_count_toward_overall_but_neither_group():
_run("k", {0: 1.0}, 0.0)
p = bi.profile("k")
assert p["n_liked"] == 0 and p["n_disliked"] == 0
assert p["overall"][0] == pytest.approx(1.0)
# --- flatness: the number the whole probe exists to produce ---------------------
def test_a_perfectly_flat_profile_reports_zero_flatness():
"""Every block moving the stream by the same amount is the null result that says
rating-driven block weighting has nothing to grip."""
_run("k", {b: 0.5 for b in range(10)}, 1.0)
assert bi.profile("k")["flatness"] == pytest.approx(0.0)
def test_a_structured_profile_reports_nonzero_flatness():
_run("k", {0: 0.1, 1: 0.5, 2: 1.5}, 1.0)
assert bi.profile("k")["flatness"] > 0.5
def test_flatness_is_scale_invariant():
"""Coefficient of variation, not raw spread: doubling every block's delta describes the
same shape and must not read as twice as structured."""
_run("a", {0: 0.1, 1: 0.2, 2: 0.4}, 1.0)
_run("b", {0: 0.2, 1: 0.4, 2: 0.8}, 1.0)
assert bi.profile("a")["flatness"] == pytest.approx(bi.profile("b")["flatness"])
def test_flatness_is_none_with_no_data():
assert bi.profile("k")["flatness"] is None
def test_corrupt_state_file_degrades_to_empty(tmp_path, monkeypatch):
path = tmp_path / "junk.pt"
path.write_bytes(b"not a torch file")
monkeypatch.setattr(bi, "state_path", lambda key: str(path))
assert bi.profile("k")["overall"] == {}
assert bi.commit("k", 1.0) == "no_pending"
def test_rows_survive_a_reload_as_plain_floats():
"""weights_only=True load: the state must never depend on unpickling a custom class."""
_run("k", {0: 0.25}, 1.0)
rows = bi._load("k")["rows"]
assert isinstance(rows[0]["profile"][0], float)
assert isinstance(rows[0]["weight"], float)
assert not isinstance(rows[0]["profile"][0], torch.Tensor)
# --- the sampler hook itself ---------------------------------------------------
#
# The module above is only half the probe. The half that can silently record nothing is the
# hook: if the video mask misses, the input tensor is not where it is expected, or the
# swallow-everything guard eats a real error, the profile comes back empty and looks like
# "this model has no block structure" rather than "the probe never ran". These exercise the
# registered hook directly, the way the model calls it.
import types # noqa: E402
from pathlib import Path # noqa: E402
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
for _name in (
"comfy", "comfy.k_diffusion", "comfy.k_diffusion.sampling",
"comfy.model_sampling", "comfy.nested_tensor", "comfy.sample",
"comfy.samplers", "comfy.utils",
):
sys.modules.setdefault(_name, types.ModuleType(_name))
sys.modules["comfy.nested_tensor"].NestedTensor = object
import samplers # noqa: E402
class _Model:
def __init__(self):
self.model_options = {}
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 _install(capture, n_blocks=2, max_rows=512):
s = samplers.FunPackLTXAVSceneChainSampler()
patched = s._install_block_influence(_Model(), capture, n_blocks=n_blocks,
max_rows=max_rows)
return patched.model_options["transformer_options"]["patches_replace"]["dit"]
def _ratio(entries):
"""capture[0][block] holds [||delta||, ||before||] pairs -- the ratio is derived, not
stored, so the readout can normalise by something comparable across depth."""
m = torch.stack(entries).mean(dim=0)
return float(m[0] / m[1])
def _call(hook, src, out, mod_segments):
"""Drive the hook the way minimax/model.py does."""
return hook({"img": src, "mod_segments": mod_segments},
{"original_block": lambda a: {"img": out}})
def test_hook_records_the_relative_delta_on_video_rows():
capture = [{}]
dit = _install(capture)
src = torch.zeros(4, 8)
src[:, 0] = 3.0 # ||video rows|| = 6.0 over 4 rows
out = src.clone()
out[:, 1] = 3.0 # delta has the same norm as the input
res = _call(dit[("double_block", 0)], src, out, [(0, 4, 6)]) # tag 6 % 3 == 0 -> video
assert torch.equal(res["img"], out) # measurement must not modify the stream
recorded = _ratio(capture[0][0])
assert recorded == pytest.approx(1.0, rel=1e-4)
def test_hook_averages_across_steps():
capture = [{}]
dit = _install(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
for scale in (1.0, 3.0): # two "steps" with different deltas
out = src.clone()
out[:, 1] = scale
_call(dit[("double_block", 0)], src, out, [(0, 4, 6)])
assert _ratio(capture[0][0]) == pytest.approx(2.0, rel=1e-4)
def test_text_and_audio_rows_are_not_measured():
"""Only video rows count -- a huge delta on text rows must not register."""
capture = [{}]
dit = _install(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
out = src.clone()
out[2:, 1] = 50.0 # rows 2-3 are text below
_call(dit[("double_block", 0)], src, out, [(0, 2, 6), (2, 4, 7)])
recorded = _ratio(capture[0][0])
assert recorded == pytest.approx(0.0, abs=1e-6)
def test_no_video_segment_records_nothing_rather_than_zero():
capture = [{}]
dit = _install(capture)
src = torch.ones(4, 8)
_call(dit[("double_block", 0)], src, src * 2, [(0, 4, 7)]) # text only
assert capture[0] == {}
def test_each_block_accumulates_separately():
capture = [{}]
dit = _install(capture, n_blocks=3)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
for block, scale in ((0, 1.0), (2, 4.0)):
out = src.clone()
out[:, 1] = scale
_call(dit[("double_block", block)], src, out, [(0, 4, 6)])
assert set(capture[0]) == {0, 2} # block 1 was never called
assert _ratio(capture[0][0]) == pytest.approx(1.0, rel=1e-4)
assert _ratio(capture[0][2]) == pytest.approx(4.0, rel=1e-4)
def test_row_subsampling_bounds_the_reduction():
"""The stride keeps the cost flat as the sequence grows; the estimate must stay close."""
capture = [{}]
dit = _install(capture, max_rows=16)
n = 4096
src = torch.zeros(n, 8)
src[:, 0] = 1.0
out = src.clone()
out[:, 1] = 2.0
_call(dit[("double_block", 0)], src, out, [(0, n, 6)])
assert _ratio(capture[0][0]) == pytest.approx(2.0, rel=1e-3)
def test_an_existing_patch_at_the_same_block_is_chained_not_replaced():
"""Installing over another mechanism's hook must call it, not drop it."""
calls = []
class _M(_Model):
pass
s = samplers.FunPackLTXAVSceneChainSampler()
base = _M()
def _inner(args, extra):
calls.append(True)
return {"img": args["img"] * 2}
base.model_options["transformer_options"] = {
"patches_replace": {"dit": {("double_block", 0): _inner}}}
capture = [{}]
patched = s._install_block_influence(base, capture, n_blocks=1)
hook = patched.model_options["transformer_options"]["patches_replace"]["dit"][
("double_block", 0)]
src = torch.ones(4, 8)
res = _call(hook, src, src, [(0, 4, 6)])
assert calls == [True] # inner ran
assert torch.equal(res["img"], src * 2) # and its output is what propagates
def test_a_broken_hook_never_breaks_sampling():
capture = [{}]
dit = _install(capture)
src = torch.ones(4, 8)
# mod_segments of the wrong shape entirely: must pass the block output through anyway.
res = _call(dit[("double_block", 0)], src, src * 3, "not-segments")
assert torch.equal(res["img"], src * 3)
assert capture[0] == {}
# --- the collection switch ------------------------------------------------------
#
# Research data is opt-in from Settings > Refinement & Taste. This switch gates COLLECTION
# only -- it must never be wired to anything that records a rating (see the 2026-09-03
# session: a "pause capture" toggle that silently stopped writing ratings).
@pytest.fixture
def _switch(monkeypatch, tmp_path):
monkeypatch.delenv(bi._ENV_SWITCH, raising=False)
monkeypatch.setattr(bi, "_switch_dir", lambda: str(tmp_path / "sw"))
return tmp_path
def test_collection_is_off_by_default(_switch):
assert bi.collection_enabled() is False
def test_setting_it_on_persists_to_disk_and_to_this_process(_switch, monkeypatch):
assert bi.set_collection_enabled(True) is True
assert bi.collection_enabled() is True
# A restarted process has no env var; the on-disk copy must still say on.
monkeypatch.delenv(bi._ENV_SWITCH, raising=False)
assert bi.collection_enabled() is True
def test_turning_it_off_persists_too(_switch, monkeypatch):
bi.set_collection_enabled(True)
bi.set_collection_enabled(False)
monkeypatch.delenv(bi._ENV_SWITCH, raising=False)
assert bi.collection_enabled() is False
def test_env_var_wins_over_disk(_switch, monkeypatch):
bi.set_collection_enabled(False)
monkeypatch.setenv(bi._ENV_SWITCH, "1")
assert bi.collection_enabled() is True
def test_unwritable_switch_dir_still_sets_this_session(_switch, monkeypatch):
monkeypatch.setattr(bi, "_switch_dir", lambda: "/nonexistent-root/nope")
assert bi.set_collection_enabled(True) is True
assert bi.collection_enabled() is True # env var carries it for this process
def test_the_switch_does_not_gate_recording(_switch, monkeypatch, tmp_path):
"""HARD RULE: this toggle controls whether the PROBE is installed, nothing else. It must
never stand between a rating and the data it belongs to -- save_pending/commit are
unconditional, exactly as h3_repr_steering's are."""
monkeypatch.setattr(bi, "state_path", lambda key: str(tmp_path / f"{key}.pt"))
bi.set_collection_enabled(False)
bi.save_pending("k", {0: 0.5})
assert bi.commit("k", 1.0) == "recorded"
assert bi.profile("k")["overall"][0] == pytest.approx(0.5)
# --- novelty: what each block ADDS that the one before it didn't ------------------
#
# Magnitude alone cannot tell "this block added something new" from "this block pushed
# harder in the direction the last one already went". The cosine with the previous block's
# delta is what separates those, and it is the only number here that speaks to whether the
# 50 blocks are doing 50 different things or one thing 50 times.
def test_novelty_is_recorded_and_averaged():
bi.save_pending("k", {1: 0.1, 2: 0.2}, novelty={2: 0.0})
bi.commit("k", 1.0)
bi.save_pending("k", {1: 0.1, 2: 0.2}, novelty={2: 1.0})
bi.commit("k", 1.0)
p = bi.profile("k")
assert p["novelty"][2] == pytest.approx(0.5)
assert p["mean_novelty"] == pytest.approx(0.5)
def test_novelty_is_optional_and_old_rows_stay_readable():
"""A row recorded before novelty existed carries none; it must not break the readout or
be counted as a 0.0 cosine (which would read as 'perfectly orthogonal')."""
bi.save_pending("k", {1: 0.5}) # no novelty
bi.commit("k", 1.0)
bi.save_pending("k", {1: 0.5}, novelty={1: 0.8})
bi.commit("k", 1.0)
p = bi.profile("k")
assert p["novelty"][1] == pytest.approx(0.8) # averaged over the row that HAS it only
def test_no_novelty_anywhere_reports_none_not_zero():
bi.save_pending("k", {1: 0.5})
bi.commit("k", 1.0)
p = bi.profile("k")
assert p["novelty"] == {}
assert p["mean_novelty"] is None
def test_negative_novelty_survives_the_round_trip():
"""Blocks partly undoing each other is a real outcome and the most interesting one --
it must not be clamped or dropped on the way through."""
bi.save_pending("k", {3: 0.4}, novelty={3: -0.6})
bi.commit("k", 1.0)
assert bi.profile("k")["novelty"][3] == pytest.approx(-0.6)
# --- the hook's novelty half ------------------------------------------------------
def _install2(capture, n_blocks=4, max_rows=512):
s = samplers.FunPackLTXAVSceneChainSampler()
patched = s._install_block_influence(_Model(), capture, n_blocks=n_blocks,
max_rows=max_rows)
return patched.model_options["transformer_options"]["patches_replace"]["dit"]
def _push(dit, block, src, delta_vec, segs=None):
out = src + delta_vec
return _call(dit[("double_block", block)], src, out, segs or [(0, 4, 6)])
def test_hook_records_orthogonal_deltas_as_zero_novelty():
capture = [{}, {}]
dit = _install2(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
d0 = torch.zeros(4, 8); d0[:, 1] = 1.0
d1 = torch.zeros(4, 8); d1[:, 2] = 1.0 # orthogonal to d0
_push(dit, 0, src, d0)
_push(dit, 1, src, d1)
assert torch.stack(capture[1][1]).mean().item() == pytest.approx(0.0, abs=1e-5)
def test_hook_records_a_repeated_direction_as_novelty_one():
capture = [{}, {}]
dit = _install2(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
d = torch.zeros(4, 8); d[:, 1] = 1.0
_push(dit, 0, src, d)
_push(dit, 1, src, d * 3.0) # same direction, bigger
assert torch.stack(capture[1][1]).mean().item() == pytest.approx(1.0, rel=1e-4)
def test_hook_records_an_undoing_block_as_negative_novelty():
capture = [{}, {}]
dit = _install2(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
d = torch.zeros(4, 8); d[:, 1] = 1.0
_push(dit, 0, src, d)
_push(dit, 1, src, -d)
assert torch.stack(capture[1][1]).mean().item() == pytest.approx(-1.0, rel=1e-4)
def test_the_first_block_of_a_step_has_no_predecessor():
"""Block 0 of step 2 must not be compared against block 49 of step 1 -- that pairs the
end of one denoise step with the start of the next, which is not a depth relationship."""
capture = [{}, {}]
dit = _install2(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
d = torch.zeros(4, 8); d[:, 1] = 1.0
for _step in range(2): # two steps, blocks 0..2 each
for b in range(3):
_push(dit, b, src, d)
assert 0 not in capture[1] # never a novelty entry for the first block
assert len(capture[1][1]) == 2 # one per step, not three
def test_novelty_is_skipped_when_the_holder_has_no_slot_for_it():
"""Back-compat: a caller passing a 1-slot holder still gets magnitudes and no crash."""
capture = [{}]
dit = _install2(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
d = torch.zeros(4, 8); d[:, 1] = 1.0
_push(dit, 0, src, d)
_push(dit, 1, src, d)
assert set(capture[0]) == {0, 1}
assert len(capture) == 1
def test_a_block_that_writes_in_place_is_still_measured():
"""THE regression that shipped: `before` was read AFTER the block ran, so a block writing
its residual in place left src already mutated and every delta came out exactly zero --
which then made every novelty cosine zero too. The whole readout was 0.0000 / +0.00 and
read as "this model has no block structure" rather than "the probe measured nothing"."""
capture = [{}, {}]
dit = _install(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
def in_place(args):
args["img"][:, 1] += 2.0 # writes into the caller's tensor, returns it
return {"img": args["img"]}
res = dit[("double_block", 0)]({"img": src, "mod_segments": [(0, 4, 6)]},
{"original_block": in_place})
recorded = _ratio(capture[0][0])
assert recorded == pytest.approx(2.0, rel=1e-4) # |delta| = 2*|before|
assert torch.equal(res["img"], src)
def test_a_block_returning_a_new_tensor_is_unaffected_by_the_fix():
capture = [{}, {}]
dit = _install(capture)
src = torch.zeros(4, 8)
src[:, 0] = 1.0
out = src.clone()
out[:, 1] = 3.0
_call(dit[("double_block", 0)], src, out, [(0, 4, 6)])
assert _ratio(capture[0][0]) == pytest.approx(3.0, rel=1e-4)
# --- share vs ratio: the depth confound -------------------------------------------
#
# ||f_i||/||x_i|| is not comparable across depth. The residual stream GROWS as it goes, so
# an early block scores high because its denominator is barely formed, not because it moves
# the picture more. Live data made this concrete: block 1 read 6.01 and block 36 read 0.368,
# a 16x gap that is mostly the stream norm. `share` divides every block by the same total
# instead, so the ranking answers "who moved the picture" rather than "who ran early".
def test_share_is_each_blocks_fraction_of_all_movement():
bi.save_pending("k", {0: 9.0, 1: 1.0}, raw={0: 1.0, 1: 3.0})
bi.commit("k", 1.0)
p = bi.profile("k")
assert p["share"][0] == pytest.approx(0.25)
assert p["share"][1] == pytest.approx(0.75)
# ...and the ratio still says the opposite, which is exactly the confound.
assert p["overall"][0] > p["overall"][1]
def test_share_is_absent_rather_than_wrong_without_raw_norms():
bi.save_pending("k", {0: 0.5})
bi.commit("k", 1.0)
assert bi.profile("k")["share"] == {}
def test_share_averages_across_runs_before_normalising():
for raw in ({0: 1.0, 1: 1.0}, {0: 3.0, 1: 1.0}):
bi.save_pending("k", {0: 0.1, 1: 0.1}, raw=raw)
bi.commit("k", 1.0)
p = bi.profile("k")
assert p["raw"][0] == pytest.approx(2.0) # mean of 1 and 3
assert p["share"][0] == pytest.approx(2.0 / 3.0)