digital-garbageandClaude Opus 5
3a62759be9
Settings ▸ Learning ▸ "Start a fresh measurement". Deleting the refinement key was not enough: the probe log lives outside the <key>.* sidecar convention on purpose, so it survives a key clear, and the per-bucket heads are REBUILT from it on every rating — the next rating restored them from runs still on disk. So the clear takes both, and leaves the switch alone: "throw away what was measured" and "stop measuring" are different requests. Verified against a running ComfyUI: 12 runs and the heads file removed, count back to 0. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
807 lines
35 KiB
Python
807 lines
35 KiB
Python
"""Trajectory probe (trajectory_probe.py) + its sampler observer.
|
|
|
|
The probe exists to answer one question — do ratings separate in x0_hat space at EARLY
|
|
schedule positions, where nothing currently steers — so these tests cover the two things
|
|
that would make its answer wrong rather than merely absent:
|
|
|
|
* recording: H3 bucketing over shift-generated schedules, the clamp to few-step runs,
|
|
the CFG>1 dedupe, per-scene separation, and the run/rating pairing;
|
|
* the statistic: separable data has to register, noise must NOT, and a separation that
|
|
is really about which prompt was running must not pass as a rating effect.
|
|
|
|
Plus the observer's one hard contract: it records without changing the prediction.
|
|
"""
|
|
import os
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
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
|
|
import trajectory_probe as probe # noqa: E402
|
|
|
|
|
|
def _patch_store(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(probe, "_probe_dir", lambda: str(tmp_path))
|
|
|
|
|
|
def _h3_sigmas(steps, shift=12.0):
|
|
"""H3's schedule: sigma = shift*t / (1 + (shift-1)*t) over uniform t, descending.
|
|
|
|
The point of bucketing by POSITION is that this stays near 1.0 for most of the run at
|
|
shift 12 — the exact reason the steering gate had to stop reading absolute sigma.
|
|
"""
|
|
out = []
|
|
for i in range(steps + 1):
|
|
t = 1.0 - i / steps
|
|
out.append(shift * t / (1.0 + (shift - 1.0) * t))
|
|
return torch.tensor(out)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bucketing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_h3_buckets_follow_position_not_sigma_value():
|
|
"""On a 12-step shift-12 schedule every sigma but the last sits above 0.5, so an
|
|
absolute-sigma reading would put the whole run in one bucket. Position must not."""
|
|
sigmas = _h3_sigmas(12, shift=12.0)
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=4)
|
|
buckets = [rec.bucket_for(float(sigmas[k])) for k in range(12)]
|
|
assert buckets == [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]
|
|
# ...and the naive reading really would have collapsed: most steps are still high sigma.
|
|
assert sum(1 for k in range(12) if float(sigmas[k]) > 0.5) >= 9
|
|
|
|
|
|
def test_buckets_clamp_to_step_count_on_turbo():
|
|
"""4-step turbo asked for 4 buckets gets one step each; a 3-step run cannot report 4
|
|
buckets it never measured."""
|
|
assert probe.TrajectoryRecorder(_h3_sigmas(4), n_buckets=4).n_buckets == 4
|
|
assert probe.TrajectoryRecorder(_h3_sigmas(3), n_buckets=4).n_buckets == 3
|
|
assert probe.TrajectoryRecorder(_h3_sigmas(7), n_buckets=4).n_buckets == 4
|
|
|
|
|
|
def test_recorder_without_a_schedule_records_nothing():
|
|
rec = probe.TrajectoryRecorder(None)
|
|
assert rec.bucket_for(0.5) is None
|
|
assert rec.observe(0.5, torch.ones(probe.DESCRIPTOR_DIM)) is False
|
|
assert rec.is_empty()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Recording
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_bucket_mean_averages_its_steps():
|
|
sigmas = _h3_sigmas(4)
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=2)
|
|
rec.observe(float(sigmas[0]), torch.full((probe.DESCRIPTOR_DIM,), 1.0))
|
|
rec.observe(float(sigmas[1]), torch.full((probe.DESCRIPTOR_DIM,), 3.0))
|
|
rows = rec.cell_rows()
|
|
assert [r["bucket"] for r in rows] == [0]
|
|
assert rows[0]["steps"] == 2
|
|
assert torch.allclose(rows[0]["desc"].float(), torch.full((probe.DESCRIPTOR_DIM,), 2.0))
|
|
|
|
|
|
def test_repeated_sigma_is_recorded_once():
|
|
"""H3 samples at CFG 1.0 (one call per step), so this never fires there. On a CFG>1
|
|
model it stops the unconditional prediction being averaged into the measurement."""
|
|
sigmas = _h3_sigmas(4)
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=4)
|
|
assert rec.observe(float(sigmas[0]), torch.ones(probe.DESCRIPTOR_DIM)) is True
|
|
assert rec.observe(float(sigmas[0]), torch.zeros(probe.DESCRIPTOR_DIM)) is False
|
|
rows = rec.cell_rows()
|
|
assert rows[0]["steps"] == 1
|
|
assert torch.allclose(rows[0]["desc"].float(), torch.ones(probe.DESCRIPTOR_DIM))
|
|
|
|
|
|
def test_scenes_do_not_average_into_each_other():
|
|
sigmas = _h3_sigmas(4)
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=4)
|
|
rec.begin_scene(0)
|
|
rec.observe(float(sigmas[0]), torch.ones(probe.DESCRIPTOR_DIM))
|
|
rec.begin_scene(1)
|
|
rec.observe(float(sigmas[0]), torch.full((probe.DESCRIPTOR_DIM,), 5.0))
|
|
rows = rec.cell_rows()
|
|
assert [(r["scene"], r["bucket"]) for r in rows] == [(0, 0), (1, 0)]
|
|
assert float(rows[0]["desc"][0]) == 1.0 and float(rows[1]["desc"][0]) == 5.0
|
|
|
|
|
|
def test_a_second_schedules_steps_are_not_nearest_matched_into_the_first():
|
|
"""A scene can sample twice: a second pass runs its own schedule IN FULL from its own
|
|
starting sigma, with the observer still installed. Those sigmas have no position on the
|
|
first schedule, and nearest-matching them silently averages two different denoisings
|
|
into one bucket."""
|
|
rec = probe.TrajectoryRecorder(_h3_sigmas(4), n_buckets=4)
|
|
for k in range(4):
|
|
rec.observe(float(_h3_sigmas(4)[k]), torch.ones(probe.DESCRIPTOR_DIM))
|
|
first = {r["bucket"]: r["desc"].clone() for r in rec.cell_rows()}
|
|
|
|
# Every one of these is far below the first schedule's lowest step sigma, so all four
|
|
# would nearest-match into its last bucket.
|
|
for sigma in (0.35, 0.2, 0.08, 0.01):
|
|
assert rec.observe(sigma, torch.full((probe.DESCRIPTOR_DIM,), 99.0)) is False
|
|
assert rec._foreign == 4
|
|
for row in rec.cell_rows():
|
|
assert torch.equal(row["desc"], first[row["bucket"]]), "pass 1 was contaminated"
|
|
|
|
|
|
def test_a_bound_second_pass_is_measured_separately():
|
|
second = torch.tensor([0.35, 0.2, 0.08, 0.0])
|
|
rec = probe.TrajectoryRecorder(_h3_sigmas(4), n_buckets=4)
|
|
for k in range(4):
|
|
rec.observe(float(_h3_sigmas(4)[k]), torch.ones(probe.DESCRIPTOR_DIM))
|
|
rec.begin_pass(second, index=1)
|
|
for k in range(3):
|
|
rec.observe(float(second[k]), torch.full((probe.DESCRIPTOR_DIM,), 99.0))
|
|
|
|
rows = rec.cell_rows()
|
|
assert {r["pass"] for r in rows} == {0, 1}
|
|
assert all(float(r["desc"][0]) == 1.0 for r in rows if r["pass"] == 0)
|
|
assert all(float(r["desc"][0]) == 99.0 for r in rows if r["pass"] == 1)
|
|
assert rec._foreign == 0
|
|
assert rec.pass_steps == {0: 4, 1: 3}
|
|
|
|
|
|
def test_analysis_reads_the_first_pass_only():
|
|
"""Pass 2's bucket 1 is a different sigma window than pass 1's, so pooling them would
|
|
compare unlike things while reporting a single bucket."""
|
|
rows = [{
|
|
"reward": 1.0, "prompt_hash": "p",
|
|
"rows": [{"bucket": 0, "scene": 0, "pass": 0, "desc": torch.ones(4)},
|
|
{"bucket": 0, "scene": 0, "pass": 1, "desc": torch.full((4,), 50.0)}],
|
|
}, {
|
|
"reward": -1.0, "prompt_hash": "p",
|
|
"rows": [{"bucket": 0, "scene": 0, "pass": 0, "desc": torch.zeros(4)},
|
|
{"bucket": 0, "scene": 0, "pass": 1, "desc": torch.full((4,), 50.0)}],
|
|
}]
|
|
descriptors, labels, _ = probe.collect_bucket(rows, 0)
|
|
assert labels == [1, 0]
|
|
assert torch.equal(descriptors[0], torch.ones(4))
|
|
assert torch.equal(descriptors[1], torch.zeros(4))
|
|
second, _, _ = probe.collect_bucket(rows, 0, pass_index=1)
|
|
assert all(torch.equal(d, torch.full((4,), 50.0)) for d in second)
|
|
|
|
|
|
def test_unbound_steps_are_recorded_on_the_run(monkeypatch, tmp_path):
|
|
"""A measurement taken from fewer steps than the run actually sampled has to say so."""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
rec = probe.TrajectoryRecorder(_h3_sigmas(4), n_buckets=4)
|
|
for k in range(4):
|
|
rec.observe(float(_h3_sigmas(4)[k]), torch.ones(probe.DESCRIPTOR_DIM))
|
|
rec.observe(0.123456, torch.ones(probe.DESCRIPTOR_DIM))
|
|
probe.save_pending("k", rec)
|
|
probe.commit("k", 1.0)
|
|
assert probe.load_log("k")[0]["unbound_steps"] == 1
|
|
|
|
|
|
def test_wrong_width_descriptor_is_refused():
|
|
rec = probe.TrajectoryRecorder(_h3_sigmas(4))
|
|
assert rec.observe(1.0, torch.ones(8)) is False
|
|
assert rec.is_empty()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Run/rating pairing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _recorded(sigmas, value=1.0, n_buckets=4):
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=n_buckets)
|
|
for k in range(len(sigmas) - 1):
|
|
rec.observe(float(sigmas[k]), torch.full((probe.DESCRIPTOR_DIM,), value))
|
|
return rec
|
|
|
|
|
|
def test_pending_pairs_with_the_rating_that_scores_it(monkeypatch, tmp_path):
|
|
_patch_store(monkeypatch, tmp_path)
|
|
sigmas = _h3_sigmas(4)
|
|
assert probe.save_pending("k", _recorded(sigmas), prompt_hash="p1", seed=7,
|
|
meta={"h3": True, "shift_video": 12.0, "shift_audio": 3.0})
|
|
assert probe.commit("k", 1.0, rating_key="perfect") == 1
|
|
|
|
rows = probe.load_log("k")
|
|
assert len(rows) == 1
|
|
assert rows[0]["reward"] == 1.0 and rows[0]["rating"] == "perfect"
|
|
assert rows[0]["prompt_hash"] == "p1" and rows[0]["seed"] == 7
|
|
assert rows[0]["steps"] == 4
|
|
assert probe.schedule_id(rows[0]) == (4, 12.0, 3.0)
|
|
# The pending is consumed, so a later unrelated rating cannot log this run twice.
|
|
assert probe.commit("k", -1.0) is None
|
|
assert len(probe.load_log("k")) == 1
|
|
|
|
|
|
def test_commit_without_a_run_logs_nothing(monkeypatch, tmp_path):
|
|
_patch_store(monkeypatch, tmp_path)
|
|
assert probe.commit("k", 1.0) is None
|
|
assert probe.load_log("k") == []
|
|
|
|
|
|
def test_a_run_that_measured_nothing_still_replaces_the_pending(monkeypatch, tmp_path):
|
|
"""Otherwise the previous run's pending survives and the NEXT rating consumes it, filing
|
|
one run's descriptors under another run's rating — a silently mislabelled row in the one
|
|
dataset the module exists to produce."""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
sigmas = _h3_sigmas(4)
|
|
probe.save_pending("k", _recorded(sigmas, value=1.0), prompt_hash="run-A", seed=111)
|
|
|
|
# Run B measures nothing (every step refused — e.g. a CFG>1 batch on every call).
|
|
empty = probe.TrajectoryRecorder(sigmas)
|
|
empty.note_unmeasurable()
|
|
assert empty.is_empty()
|
|
probe.save_pending("k", empty, prompt_hash="run-B", seed=222)
|
|
|
|
# Rating run B must not log run A's descriptors under run B's rating.
|
|
assert probe.commit("k", 1.0, rating_key="perfect") is None
|
|
assert probe.load_log("k") == []
|
|
# ...and the stale candidate is gone, so no later rating can pick it up either.
|
|
assert probe.commit("k", -1.0) is None
|
|
|
|
|
|
def test_save_pending_needs_a_key_and_a_recorder(monkeypatch, tmp_path):
|
|
_patch_store(monkeypatch, tmp_path)
|
|
assert probe.save_pending(None, _recorded(_h3_sigmas(4))) is False
|
|
assert probe.save_pending("k", None) is False
|
|
|
|
|
|
def test_log_is_capped(monkeypatch, tmp_path):
|
|
_patch_store(monkeypatch, tmp_path)
|
|
monkeypatch.setattr(probe, "MAX_ROWS", 3)
|
|
sigmas = _h3_sigmas(4)
|
|
for i in range(5):
|
|
probe.save_pending("k", _recorded(sigmas, value=float(i)))
|
|
probe.commit("k", float(i))
|
|
rows = probe.load_log("k")
|
|
assert len(rows) == 3
|
|
assert [r["reward"] for r in rows] == [2.0, 3.0, 4.0] # oldest rolled off
|
|
|
|
|
|
def test_logs_pool_across_disposable_keys(monkeypatch, tmp_path):
|
|
"""Keys get cleared between rentals; the measurement is supposed to outlive them."""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
for key in ("alpha", "beta"):
|
|
probe.save_pending(key, _recorded(_h3_sigmas(4)))
|
|
probe.commit(key, 1.0)
|
|
assert sorted(k for k, _ in probe.load_all_logs()) == ["alpha", "beta"]
|
|
assert sum(len(r) for _, r in probe.load_all_logs()) == 2
|
|
|
|
|
|
def test_schedule_mix_names_incompatible_runs():
|
|
rows = [
|
|
{"steps": 4, "meta": {"shift_video": 12.0, "shift_audio": 3.0}},
|
|
{"steps": 4, "meta": {"shift_video": 12.0, "shift_audio": 3.0}},
|
|
{"steps": 12, "meta": {"shift_video": 3.0, "shift_audio": 3.0}},
|
|
]
|
|
mix = probe.schedule_mix(rows)
|
|
assert mix[(4, 12.0, 3.0)] == 2 and mix[(12, 3.0, 3.0)] == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The statistic
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _descriptors(offsets, dim=32, noise=0.05):
|
|
gen = torch.Generator().manual_seed(0)
|
|
return [torch.randn(dim, generator=gen) * noise + off for off in offsets]
|
|
|
|
|
|
def test_separable_groups_register():
|
|
base = torch.zeros(32)
|
|
shift = torch.zeros(32)
|
|
shift[0] = 4.0
|
|
descriptors = _descriptors([base] * 5 + [shift] * 5)
|
|
labels = [1] * 5 + [0] * 5
|
|
result = probe.permutation_test(descriptors, labels, trials=500, seed=1)
|
|
assert result["separation"] > 0
|
|
assert result["p_value"] <= 0.05
|
|
|
|
|
|
def test_noise_does_not_register():
|
|
gen = torch.Generator().manual_seed(3)
|
|
descriptors = [torch.randn(32, generator=gen) for _ in range(12)]
|
|
labels = [1, 0] * 6
|
|
result = probe.permutation_test(descriptors, labels, trials=500, seed=1)
|
|
assert result["p_value"] > 0.05
|
|
|
|
|
|
def test_prompt_confound_does_not_pass_as_a_rating_effect():
|
|
"""Every good run on prompt A, every bad run on prompt B. Pooled, that separates
|
|
perfectly and means nothing. Restricted to same-prompt pairs there is no evidence at
|
|
all — which is why the analysis reports the stratified row."""
|
|
a, b = torch.zeros(32), torch.zeros(32)
|
|
b[0] = 6.0
|
|
descriptors = _descriptors([a] * 4 + [b] * 4)
|
|
labels = [1] * 4 + [0] * 4
|
|
groups = ["A"] * 4 + ["B"] * 4
|
|
|
|
assert probe.permutation_test(descriptors, labels, trials=300, seed=1)["p_value"] <= 0.05
|
|
# No pair inside a prompt has differing labels, so there is nothing to compare.
|
|
assert probe.separation_statistic(descriptors, labels, groups) is None
|
|
assert probe.permutation_test(descriptors, labels, groups, trials=300, seed=1) is None
|
|
|
|
|
|
def test_statistic_needs_both_labels():
|
|
descriptors = _descriptors([torch.zeros(32)] * 4)
|
|
assert probe.separation_statistic(descriptors, [1, 1, 1, 1]) is None
|
|
assert probe.separation_statistic(descriptors[:2], [1, 0]) is None
|
|
|
|
|
|
def test_neutral_ratings_are_dropped_not_assigned():
|
|
"""'Forget it' carries reward 0.0 and is explicitly not a judgement."""
|
|
rows = [
|
|
{"reward": 1.0, "prompt_hash": "p", "rows": [{"bucket": 0, "scene": 0, "desc": torch.ones(4)}]},
|
|
{"reward": 0.0, "prompt_hash": "p", "rows": [{"bucket": 0, "scene": 0, "desc": torch.ones(4)}]},
|
|
{"reward": -0.5, "prompt_hash": "p", "rows": [{"bucket": 0, "scene": 0, "desc": torch.zeros(4)}]},
|
|
]
|
|
descriptors, labels, _ = probe.collect_bucket(rows, 0)
|
|
assert labels == [1, 0] and len(descriptors) == 2
|
|
|
|
|
|
def test_analyse_reports_every_recorded_bucket():
|
|
def _row(reward, scale):
|
|
return {"reward": reward, "prompt_hash": "p",
|
|
"rows": [{"bucket": b, "scene": 0, "desc": torch.full((8,), scale * (b + 1))}
|
|
for b in range(4)]}
|
|
rows = [_row(1.0, 1.0) for _ in range(3)] + [_row(-1.0, -1.0) for _ in range(3)]
|
|
results = probe.analyse(rows, trials=100)
|
|
assert [e["bucket"] for e in results] == [0, 1, 2, 3]
|
|
assert all(e["n_good"] == 3 and e["n_bad"] == 3 for e in results)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The sampler observer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
C, T, H, W = 4, 3, 2, 2
|
|
VIDEO_N = C * T * H * W
|
|
AUDIO_N = 12
|
|
|
|
|
|
def _packed_av_model():
|
|
ls = types.SimpleNamespace(cond=[(1, C, T, H, W), (1, AUDIO_N)])
|
|
guider = types.SimpleNamespace(conds={"positive": [{"model_conds": {"latent_shapes": ls}}]})
|
|
return types.SimpleNamespace(model_options={}, inner_model=guider)
|
|
|
|
|
|
def test_observer_records_without_touching_the_prediction():
|
|
model = _packed_av_model()
|
|
sigmas = _h3_sigmas(4)
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=4)
|
|
sampler = samplers.FunPackLTXAVSceneChainSampler()
|
|
sampler._build_trajectory_probe_wrapper(model, rec)
|
|
wrapper = model.model_options["model_function_wrapper"]
|
|
|
|
denoised = torch.arange(VIDEO_N + AUDIO_N, dtype=torch.float32).view(1, 1, -1)
|
|
for k in range(4):
|
|
out = wrapper(lambda x, t, **kw: denoised,
|
|
{"input": denoised, "timestep": torch.tensor([float(sigmas[k])]), "c": {}})
|
|
assert torch.equal(out, denoised), "the probe must observe, never steer"
|
|
assert [r["bucket"] for r in rec.cell_rows()] == [0, 1, 2, 3]
|
|
|
|
|
|
def test_observer_pools_video_only():
|
|
"""Audio must not reach the descriptor: it is not the domain LatentValueFunction is
|
|
trained on, and on H3 the audio stream is a different clock entirely."""
|
|
model = _packed_av_model()
|
|
sigmas = _h3_sigmas(4)
|
|
sampler = samplers.FunPackLTXAVSceneChainSampler()
|
|
|
|
def _descriptor_for(audio_fill):
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=4)
|
|
model.model_options.pop("model_function_wrapper", None)
|
|
sampler._build_trajectory_probe_wrapper(model, rec)
|
|
latent = torch.cat([torch.ones(VIDEO_N), torch.full((AUDIO_N,), audio_fill)]).view(1, 1, -1)
|
|
model.model_options["model_function_wrapper"](
|
|
lambda x, t, **kw: latent,
|
|
{"input": latent, "timestep": torch.tensor([float(sigmas[0])]), "c": {}})
|
|
return rec.cell_rows()[0]["desc"].float()
|
|
|
|
assert torch.allclose(_descriptor_for(0.0), _descriptor_for(99.0))
|
|
|
|
|
|
def test_descriptors_survive_inference_mode(monkeypatch, tmp_path):
|
|
"""Sampling runs under inference_mode, and these descriptors outlive the call that made
|
|
them: accumulated across steps, then written to disk. This drives the whole path that
|
|
way — record, accumulate, save, reload.
|
|
|
|
The observer takes the inference_mode(False) precaution _save_output_value_snapshot
|
|
takes, but this test passes with or without it on the pinned torch, so read it as
|
|
coverage of the path, not as proof that the precaution is load-bearing.
|
|
"""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
model = _packed_av_model()
|
|
sigmas = _h3_sigmas(4)
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=4)
|
|
samplers.FunPackLTXAVSceneChainSampler()._build_trajectory_probe_wrapper(model, rec)
|
|
wrapper = model.model_options["model_function_wrapper"]
|
|
|
|
with torch.inference_mode():
|
|
latent = torch.arange(VIDEO_N + AUDIO_N, dtype=torch.float32).view(1, 1, -1)
|
|
for k in range(4):
|
|
wrapper(lambda x, t, **kw: latent,
|
|
{"input": latent, "timestep": torch.tensor([float(sigmas[k])]), "c": {}})
|
|
|
|
assert [r["bucket"] for r in rec.cell_rows()] == [0, 1, 2, 3]
|
|
assert probe.save_pending("k", rec, prompt_hash="p", seed=1, meta={"h3": True})
|
|
assert probe.commit("k", 1.0) == 1
|
|
assert len(probe.load_log("k")[0]["rows"]) == 4
|
|
|
|
|
|
def test_batched_prediction_is_refused_not_averaged():
|
|
"""At CFG>1 comfy stacks the unconditional row into the SAME call's batch, and the
|
|
descriptor pools over batch — recording it would average a prediction with its own
|
|
negation and report the run as cleanly measured."""
|
|
model = _packed_av_model()
|
|
sigmas = _h3_sigmas(4)
|
|
rec = probe.TrajectoryRecorder(sigmas, n_buckets=4)
|
|
samplers.FunPackLTXAVSceneChainSampler()._build_trajectory_probe_wrapper(model, rec)
|
|
wrapper = model.model_options["model_function_wrapper"]
|
|
|
|
cond = torch.full((1, 1, VIDEO_N + AUDIO_N), 10.0)
|
|
uncond = torch.zeros(1, 1, VIDEO_N + AUDIO_N)
|
|
batched = torch.cat([cond, uncond], dim=0) # comfy's cond/uncond stacking
|
|
out = wrapper(lambda x, t, **kw: batched,
|
|
{"input": batched, "timestep": torch.tensor([float(sigmas[0])]), "c": {}})
|
|
|
|
assert torch.equal(out, batched), "the probe must observe, never steer"
|
|
assert rec.is_empty(), "a cond/uncond blend must not be recorded as a measurement"
|
|
assert rec._foreign == 1, "and the run has to be able to say it was not measured"
|
|
|
|
|
|
def test_a_step_with_no_timestep_is_counted_not_dropped():
|
|
"""Every path that cannot measure a step counts it, so a run can say the measurement is
|
|
incomplete. A step with no timestep is one of those paths, not an exception to it."""
|
|
model = _packed_av_model()
|
|
rec = probe.TrajectoryRecorder(_h3_sigmas(4), n_buckets=4)
|
|
samplers.FunPackLTXAVSceneChainSampler()._build_trajectory_probe_wrapper(model, rec)
|
|
latent = torch.zeros(1, 1, VIDEO_N + AUDIO_N)
|
|
out = model.model_options["model_function_wrapper"](
|
|
lambda x, t, **kw: latent, {"input": latent, "timestep": None, "c": {}})
|
|
assert torch.equal(out, latent)
|
|
assert rec.is_empty()
|
|
assert rec._foreign == 1
|
|
|
|
|
|
def test_observer_survives_a_broken_prediction():
|
|
"""A probe failure must cost the measurement, never the generation."""
|
|
model = _packed_av_model()
|
|
rec = probe.TrajectoryRecorder(_h3_sigmas(4), n_buckets=4)
|
|
sampler = samplers.FunPackLTXAVSceneChainSampler()
|
|
sampler._build_trajectory_probe_wrapper(model, rec)
|
|
out = model.model_options["model_function_wrapper"](
|
|
lambda x, t, **kw: "not a tensor",
|
|
{"input": torch.zeros(1, 1, 4), "timestep": torch.tensor([0.9]), "c": {}})
|
|
assert out == "not a tensor"
|
|
assert rec.is_empty()
|
|
# A step lost to an exception is as unmeasured as one refused on purpose, and has to be
|
|
# counted, or a run that silently lost steps reads downstream as fully measured.
|
|
assert rec._foreign == 1
|
|
|
|
|
|
def test_observer_chains_onto_an_existing_wrapper():
|
|
"""It installs outermost, over the steering stack — whatever was already there has to
|
|
keep running, or the probe would silently disable steering it is meant to observe."""
|
|
model = _packed_av_model()
|
|
calls = []
|
|
|
|
def _existing(apply_fn, args):
|
|
calls.append("inner")
|
|
return apply_fn(args["input"], args["timestep"], **args.get("c", {}))
|
|
|
|
model.model_options["model_function_wrapper"] = _existing
|
|
rec = probe.TrajectoryRecorder(_h3_sigmas(4), n_buckets=4)
|
|
samplers.FunPackLTXAVSceneChainSampler()._build_trajectory_probe_wrapper(model, rec)
|
|
latent = torch.zeros(1, 1, VIDEO_N + AUDIO_N)
|
|
model.model_options["model_function_wrapper"](
|
|
lambda x, t, **kw: latent,
|
|
{"input": latent, "timestep": torch.tensor([0.9]), "c": {}})
|
|
assert calls == ["inner"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The rating intake (conditioning.refine_v2)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _refiner(tmp_path):
|
|
from conditioning import FunPackVideoRefinerV2
|
|
r = FunPackVideoRefinerV2()
|
|
r._v2_state_path = lambda refinement_key: str(tmp_path / "state.json")
|
|
return r
|
|
|
|
|
|
class _Clip:
|
|
def tokenize(self, text):
|
|
return text
|
|
|
|
def encode_from_tokens_scheduled(self, tokens):
|
|
return [(torch.ones(1, 4, 3), {"pooled_output": torch.ones(1, 3)})]
|
|
|
|
|
|
# The Movie Editor branch needs the run being rated to have been MULTI-SCENE, and the editor
|
|
# hands its own scene list in rather than relying on any delimiter in the prompt.
|
|
_SEGMENTS = {"anchor": "a woman", "scenes": ["walking through neon rain", "turning to camera"]}
|
|
|
|
|
|
def _rate(refiner, key, rating, me_ratings=None, segments=_SEGMENTS):
|
|
return refiner.refine_v2("a woman walking through neon rain", _Clip(), rating, key,
|
|
split_by_transitions=True, scene_segments=segments,
|
|
movie_editor_scene_ratings=me_ratings)
|
|
|
|
|
|
def test_a_rating_logs_the_run_whichever_ui_rated_it(monkeypatch, tmp_path):
|
|
"""Multi-scene chains rated per scene from the Movie Editor go down a different branch
|
|
of refine_v2 than a single overall rating. A run rated there is a labelled data point
|
|
like any other — the probe that misses it collects nothing in the workflow that
|
|
generates most of the runs."""
|
|
_patch_store(monkeypatch, tmp_path / "probe")
|
|
monkeypatch.setenv("FUNPACK_TRAJECTORY_PROBE", "1")
|
|
sigmas = _h3_sigmas(4)
|
|
|
|
refiner = _refiner(tmp_path)
|
|
_rate(refiner, "key", "Perfect") # establishes a previous run
|
|
|
|
probe.save_pending("key", _recorded(sigmas), prompt_hash="p", seed=1)
|
|
_rate(refiner, "key", "Perfect")
|
|
assert len(probe.load_log("key")) == 1, "a plain rating must log its run"
|
|
|
|
# Asserted, not assumed: the editor branch has its own preconditions (a multi-scene
|
|
# PREVIOUS run among them), and without them this test passes by quietly taking the
|
|
# single-rating path it is supposed to be contrasting with.
|
|
import conditioning
|
|
entered = []
|
|
original = conditioning.FunPackVideoRefinerV2._v2_apply_movie_editor_scene_ratings
|
|
monkeypatch.setattr(conditioning.FunPackVideoRefinerV2,
|
|
"_v2_apply_movie_editor_scene_ratings",
|
|
lambda self, *a, **kw: (entered.append(1), original(self, *a, **kw))[1])
|
|
|
|
probe.save_pending("key", _recorded(sigmas), prompt_hash="p", seed=2)
|
|
_rate(refiner, "key", "Perfect",
|
|
me_ratings=[{"index": 0, "rating": "Perfect"}, {"index": 1, "rating": "Awful"}])
|
|
# "index" is the key the editor actually sends (movie_editor/server.py). With any other
|
|
# key the aggregate lookup finds nothing and this silently becomes a second copy of the
|
|
# single-rating case above.
|
|
assert entered == [1], "the Movie Editor branch was never taken — this proves nothing"
|
|
assert len(probe.load_log("key")) == 2, "a Movie Editor per-scene rating must log too"
|
|
|
|
|
|
def test_the_intake_is_silent_while_the_probe_is_off(monkeypatch, tmp_path):
|
|
"""With the probe off the sampler writes no pending, so consuming one would pair an
|
|
older run's descriptors with this rating."""
|
|
_patch_store(monkeypatch, tmp_path / "probe")
|
|
monkeypatch.delenv("FUNPACK_TRAJECTORY_PROBE", raising=False)
|
|
|
|
refiner = _refiner(tmp_path)
|
|
_rate(refiner, "key", "Perfect")
|
|
probe.save_pending("key", _recorded(_h3_sigmas(4)))
|
|
_rate(refiner, "key", "Perfect")
|
|
|
|
assert probe.load_log("key") == []
|
|
assert os.path.exists(probe.pending_path("key")), "the pending must be left untouched"
|
|
|
|
|
|
def test_each_rated_scene_is_its_own_labelled_sample():
|
|
"""A multi-scene run rated per scene carries one reward PER SCENE. Collapsing to the
|
|
run's aggregate would file a scene the user called Perfect as a bad sample whenever some
|
|
other scene in the same run was rated badly — mixing the two labels the whole
|
|
measurement is trying to tell apart."""
|
|
row = {
|
|
"reward": -0.9, "prompt_hash": "p", # the aggregate: the worst scene wins
|
|
"scene_rewards": {0: 1.0, 1: -0.85},
|
|
"rows": [{"bucket": 0, "scene": 0, "pass": 0, "desc": torch.ones(4)},
|
|
{"bucket": 0, "scene": 1, "pass": 0, "desc": torch.zeros(4)}],
|
|
}
|
|
descriptors, labels, groups = probe.collect_bucket([row], 0)
|
|
assert labels == [1, 0], "the Perfect scene must not be labelled bad"
|
|
assert torch.equal(descriptors[0], torch.ones(4))
|
|
assert torch.equal(descriptors[1], torch.zeros(4))
|
|
assert groups == ["p", "p"]
|
|
|
|
|
|
def test_a_run_rated_once_is_still_one_sample():
|
|
row = {
|
|
"reward": 1.0, "prompt_hash": "p",
|
|
"rows": [{"bucket": 0, "scene": 0, "pass": 0, "desc": torch.ones(4)},
|
|
{"bucket": 0, "scene": 1, "pass": 0, "desc": torch.zeros(4)}],
|
|
}
|
|
descriptors, labels, _ = probe.collect_bucket([row], 0)
|
|
assert labels == [1]
|
|
assert torch.allclose(descriptors[0], torch.full((4,), 0.5)) # scenes averaged
|
|
|
|
|
|
def test_per_scene_rewards_reach_the_log(monkeypatch, tmp_path):
|
|
_patch_store(monkeypatch, tmp_path)
|
|
probe.save_pending("k", _recorded(_h3_sigmas(4)))
|
|
probe.commit("k", -0.9, rating_key="awful", scene_rewards={0: 1.0, 1: -0.85})
|
|
assert probe.load_log("k")[0]["scene_rewards"] == {0: 1.0, 1: -0.85}
|
|
|
|
|
|
def test_the_switch_survives_a_restart(monkeypatch, tmp_path):
|
|
"""A rented box restarts. If the switch lived only in the environment, recording would
|
|
stop and the panel would be the only place that ever said so."""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
monkeypatch.delenv("FUNPACK_TRAJECTORY_PROBE", raising=False)
|
|
assert probe.probe_enabled() is False
|
|
|
|
assert probe.set_probe_enabled(True) is True
|
|
assert probe.probe_enabled() is True
|
|
monkeypatch.delenv("FUNPACK_TRAJECTORY_PROBE") # the restart
|
|
assert probe.probe_enabled() is True, "the switch did not survive"
|
|
|
|
probe.set_probe_enabled(False)
|
|
monkeypatch.delenv("FUNPACK_TRAJECTORY_PROBE")
|
|
assert probe.probe_enabled() is False, "turning it off did not survive either"
|
|
|
|
|
|
def test_the_environment_wins_while_it_is_set(monkeypatch, tmp_path):
|
|
"""The sampler runs in this process, so the live toggle has to beat the saved one."""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
probe.set_probe_enabled(False)
|
|
monkeypatch.setenv("FUNPACK_TRAJECTORY_PROBE", "1")
|
|
assert probe.probe_enabled() is True
|
|
|
|
|
|
def test_the_switch_file_is_not_mistaken_for_a_log(monkeypatch, tmp_path):
|
|
_patch_store(monkeypatch, tmp_path)
|
|
probe.set_probe_enabled(True)
|
|
probe.save_pending("k", _recorded(_h3_sigmas(4)))
|
|
probe.commit("k", 1.0)
|
|
assert sorted(k for k, _ in probe.load_all_logs()) == ["k"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Carrying the measurement between machines
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_importing_the_same_export_twice_adds_nothing(monkeypatch, tmp_path):
|
|
"""A rental is replaced whenever something breaks, so an export gets carried onto a
|
|
new box and may well be loaded twice. Appending blindly would double the sample, and
|
|
the permutation test reads n directly — duplicated runs report a separation as far
|
|
more significant than the evidence supports."""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
sigmas = _h3_sigmas(4)
|
|
for i in range(3):
|
|
probe.save_pending("k", _recorded(sigmas, value=float(i)), prompt_hash="p", seed=i)
|
|
probe.commit("k", 1.0 if i else -0.9)
|
|
exported = [r for _key, rows in probe.load_all_logs() for r in rows]
|
|
assert len(exported) == 3
|
|
|
|
fresh = tmp_path / "other-box"
|
|
monkeypatch.setattr(probe, "_probe_dir", lambda: str(fresh))
|
|
assert probe.merge_rows("imported", exported) == (3, 3)
|
|
assert probe.merge_rows("imported", exported) == (3, 0)
|
|
|
|
|
|
def test_two_boxes_exports_merge(monkeypatch, tmp_path):
|
|
"""The point of carrying it: runs from separate rentals add up to one sample."""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
sigmas = _h3_sigmas(4)
|
|
|
|
def runs(stamp, n):
|
|
out = []
|
|
for i in range(n):
|
|
rec = _recorded(sigmas, value=float(i))
|
|
out.append({"rows": rec.cell_rows(), "reward": 1.0, "stamp": stamp,
|
|
"seed": i, "prompt_hash": "p"})
|
|
return out
|
|
|
|
assert probe.merge_rows("imported", runs("box-a", 3)) == (3, 3)
|
|
assert probe.merge_rows("imported", runs("box-b", 4)) == (7, 4)
|
|
|
|
|
|
def test_a_row_with_nothing_in_it_is_not_a_run(monkeypatch, tmp_path):
|
|
_patch_store(monkeypatch, tmp_path)
|
|
junk = [{}, {"rows": []}, {"rows": [{"bucket": 0}], "reward": "not a number"},
|
|
"a string", None]
|
|
assert probe.merge_rows("imported", junk) == (0, 0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Centring, and the hash
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_signal_hidden_under_shared_content_is_found():
|
|
"""Runs of one prompt share most of their content. Measured on a real session, 86% of
|
|
each descriptor was the part every run had; comparing raw vectors measured that and
|
|
reported no separation at any bucket, including the late control.
|
|
|
|
Here the shared part is 20x the difference — the regime the real data was in."""
|
|
shared = torch.ones(64) * 20.0
|
|
good = [shared + torch.full((64,), 1.0) + torch.randn(64) * 0.05 for _ in range(6)]
|
|
bad = [shared - torch.full((64,), 1.0) + torch.randn(64) * 0.05 for _ in range(6)]
|
|
labels = [1] * 6 + [0] * 6
|
|
|
|
result = probe.permutation_test(good + bad, labels, trials=400, seed=1)
|
|
assert result["separation"] > 0
|
|
assert result["p_value"] <= 0.05
|
|
|
|
|
|
def test_centring_does_not_invent_a_split_in_noise():
|
|
"""The other half: removing the shared part must not turn nothing into something."""
|
|
gen = torch.Generator().manual_seed(11)
|
|
vectors = [torch.ones(64) * 20.0 + torch.randn(64, generator=gen) for _ in range(12)]
|
|
result = probe.permutation_test(vectors, [1, 0] * 6, trials=400, seed=1)
|
|
assert result["p_value"] > 0.05
|
|
|
|
|
|
def test_one_prompt_keeps_one_id_through_gpu_jitter():
|
|
"""A real session produced 20 distinct ids across 21 generations of one prompt, and the
|
|
hash was a candidate cause: it compared raw bytes, so any value crossing a rounding
|
|
boundary changed the id.
|
|
|
|
Quantising fixes it for jitter of the size a GPU actually produces (~1e-6 relative).
|
|
It cannot be fixed for ALL jitter — with 768 values some will always sit near a
|
|
boundary, whatever the quantum — so this asserts the realistic range and not more.
|
|
"""
|
|
gen = torch.Generator().manual_seed(3)
|
|
cond = torch.randn(1, 77, 768, generator=gen)
|
|
for scale in (1e-7, 1e-6, 1e-5):
|
|
jittered = cond + torch.randn(cond.shape, generator=gen) * scale
|
|
assert probe.prompt_hash(cond) == probe.prompt_hash(jittered), scale
|
|
|
|
|
|
def test_a_different_prompt_still_gets_a_different_id():
|
|
"""The rounding must not be so coarse that two prompts collide — that would hide a
|
|
confound instead of measuring it."""
|
|
gen = torch.Generator().manual_seed(4)
|
|
a = torch.randn(1, 77, 768, generator=gen)
|
|
b = torch.randn(1, 77, 768, generator=gen)
|
|
assert probe.prompt_hash(a) != probe.prompt_hash(b)
|
|
assert probe.prompt_hash(a) == probe.prompt_hash(a.clone())
|
|
|
|
|
|
def test_clearing_removes_the_runs_and_what_was_learned(monkeypatch, tmp_path):
|
|
"""The heads are REBUILT from this log on every rating, so deleting them alone achieves
|
|
nothing — the next rating restores them from the runs still on disk. A clear that left
|
|
the log behind would look like it worked and then quietly resume steering."""
|
|
_patch_store(monkeypatch, tmp_path / "probe")
|
|
sidecars = tmp_path / "refinements"
|
|
sidecars.mkdir()
|
|
monkeypatch.setattr(probe, "_probe_dir", lambda: str(tmp_path / "probe"))
|
|
|
|
import trajectory_guidance as tg
|
|
monkeypatch.setattr(tg, "state_path", lambda key: str(sidecars / f"{key}.value_fn_buckets.pt"))
|
|
monkeypatch.setattr("conditioning.refinement_state_path",
|
|
lambda key, mode, prefix="refine", extension="json":
|
|
str(sidecars / f"{key}.{mode}.{extension}"))
|
|
|
|
for i in range(3):
|
|
probe.save_pending("k", _recorded(_h3_sigmas(4)), seed=i)
|
|
probe.commit("k", 1.0 if i else -0.9)
|
|
(sidecars / "k.value_fn_buckets.pt").write_bytes(b"heads")
|
|
assert len(probe.load_log("k")) == 3
|
|
|
|
removed = probe.clear_all()
|
|
assert removed["runs"] == 3
|
|
assert probe.load_log("k") == []
|
|
assert probe.load_all_logs() == []
|
|
assert not (sidecars / "k.value_fn_buckets.pt").exists(), "the learned windows survived"
|
|
|
|
|
|
def test_clearing_does_not_stop_recording(monkeypatch, tmp_path):
|
|
"""'Throw away what was measured' and 'stop measuring' are different requests, and a
|
|
clear that silently did both would look like the probe had broken."""
|
|
_patch_store(monkeypatch, tmp_path)
|
|
probe.set_probe_enabled(True)
|
|
probe.save_pending("k", _recorded(_h3_sigmas(4)))
|
|
probe.commit("k", 1.0)
|
|
probe.clear_all()
|
|
monkeypatch.delenv("FUNPACK_TRAJECTORY_PROBE", raising=False)
|
|
assert probe.probe_enabled() is True
|