Research data belongs in Settings > Refinement & Taste — readable, downloadable, clearable, with its own collection switch — not in a run report that already carries more than anyone but us needs. - Removed the sampler widget and its Editor knob entirely. Collection is driven by a global switch (env var live for this process + on-disk copy so a fresh rental comes back recording), the same two-layer mechanism trajectory_probe uses. - Removed the per-rating console readout. Findings go to the panel; failures still speak through funpack_log. - New section in Refinement & Taste: toggle, recorded count, Download, Clear, a plain-language flatness verdict, and the busiest blocks ranked with their liked-vs-disliked delta. - Routes: GET/POST /api/block_influence, /export, /clear. The switch gates whether the PROBE is installed and nothing else — save_pending/commit stay unconditional, with a test pinning that, since a toggle that quietly stops recording ratings is a mistake this project has already made once. Verified live: routes round-trip, the panel renders a seeded 6-run profile, and a planted rating-correlated bump at block 31 surfaces as +0.0107 in the UI. Seed data cleared afterwards.
353 lines
13 KiB
Python
353 lines
13 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 _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 = torch.stack(capture[0][0]).mean().item()
|
|
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 torch.stack(capture[0][0]).mean().item() == 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 = torch.stack(capture[0][0]).mean().item()
|
|
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 torch.stack(capture[0][0]).mean().item() == pytest.approx(1.0, rel=1e-4)
|
|
assert torch.stack(capture[0][2]).mean().item() == 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 torch.stack(capture[0][0]).mean().item() == 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)
|