diff --git a/tests/test_trajectory_guidance.py b/tests/test_trajectory_guidance.py index e1a96723..853ce486 100644 --- a/tests/test_trajectory_guidance.py +++ b/tests/test_trajectory_guidance.py @@ -62,20 +62,36 @@ def _store(monkeypatch, tmp_path): def test_a_head_learns_from_its_own_bucket_only(monkeypatch, tmp_path): """The whole claim of the feature: what steers at step 2 learned from step 2. A head - fed every bucket's descriptors would just be output_guidance with extra files.""" + fed every bucket's descriptors would just be output_guidance with extra files. + + Measured as ranking accuracy over the descriptors each head actually saw, not as a score + at some constant vector: a head trained on noise will happily extrapolate to a large + arbitrary value far outside its data, which made an earlier version of this test pass or + fail on `train_on`'s unseeded bootstrap sampling rather than on anything real. + """ + import random as _random + _random.seed(7) # train_on bootstraps with the global RNG _store(monkeypatch, tmp_path) - tg.train_from_rows("k", _separable()) + rows = _separable(n=28) + train, held_out = rows[:20], rows[20:] + tg.train_from_rows("k", train) model = tg.load("k") assert model is not None - def score(bucket, value): + def accuracy(bucket, on): + pos, neg = [], [] with torch.inference_mode(False), torch.no_grad(): - return float(model._head(bucket)(torch.full((1, DIM), value)).mean()) + for row in on: + desc = next(c["desc"] for c in row["rows"] if c["bucket"] == bucket) + score = float(model._head(bucket)(desc.unsqueeze(0)).mean()) + (pos if row["reward"] > 0 else neg).append(score) + return sum(p > q for p in pos for q in neg) / (len(pos) * len(neg)) - # Bucket 1 saw the split and must rank a good-looking descriptor above a bad one. - assert score(1, 1.0) > score(1, -1.0) - # Bucket 0 saw noise either way, so it must not have invented an opinion. - assert abs(score(0, 1.0) - score(0, -1.0)) < abs(score(1, 1.0) - score(1, -1.0)) + # HELD OUT, not the training rows. A few-hundred-parameter MLP fits 20 noise vectors + # perfectly, so training accuracy is 1.0 for every bucket and says nothing about whether + # a head learned anything -- including the head that will be steering real generations. + assert accuracy(1, held_out) > 0.9, accuracy(1, held_out) + assert accuracy(0, held_out) < 0.9, accuracy(0, held_out) def test_an_untrained_bucket_is_not_ready(monkeypatch, tmp_path): @@ -250,3 +266,40 @@ def test_it_chains_onto_whatever_is_already_installed(): wrapper = _install(model, _Ready([]), lambda s: 0) _predict(wrapper, torch.ones(1, 1, VIDEO_N + AUDIO_N)) assert calls == [1] + + +def test_training_works_inside_inference_mode(monkeypatch, tmp_path): + """ComfyUI executes a node under torch.inference_mode(), so a descriptor read out of the + log becomes an inference tensor and autograd refuses it: + + Trajectory probe intake failed: Inference tensors do not track version counter. + + Seen on a real rental. Every head stayed empty and the feature reported "nothing learned + yet" on every run, so it was OFF while appearing merely untrained — the one failure that + is invisible from the outside, because "not steering" is also what correct behaviour + looks like before enough ratings. + """ + _store(monkeypatch, tmp_path) + with torch.inference_mode(): + counts = tg.train_from_rows("k", _separable()) + assert counts, "training produced nothing under inference mode" + assert tg.load("k") is not None + + +def test_steering_works_inside_inference_mode(monkeypatch, tmp_path): + """The other half: sampling also runs under inference mode, so the gradient the wrapper + asks for is computed there too.""" + _store(monkeypatch, tmp_path) + with torch.inference_mode(): + tg.train_from_rows("k", _separable()) + model = tg.load("k") + ready = [b for b in range(model.n_buckets) if model.ready(b)] + assert ready, "nothing trained, so this would pass for the wrong reason" + + host = _model() + wrapper = _install(host, model, lambda s: ready[0]) + latent = torch.randn(1, 1, VIDEO_N + AUDIO_N) + with torch.inference_mode(): + out = _predict(wrapper, latent) + assert not torch.equal(out, latent) + assert torch.isfinite(out).all() diff --git a/trajectory_guidance.py b/trajectory_guidance.py index 4c97dc13..98040158 100644 --- a/trajectory_guidance.py +++ b/trajectory_guidance.py @@ -138,6 +138,25 @@ class BucketedValue: return out +def _train_all(model, rows): + """Every rated descriptor into the head for its own bucket.""" + for row in rows: + per_scene = row.get("scene_rewards") or None + for cell in row["rows"]: + # Pass 0 only: a second pass runs its own schedule from its own starting sigma, + # so its bucket 2 is not the window bucket 2 steers in. + if int(cell.get("pass", 0)) != 0: + continue + reward = (per_scene or {}).get(int(cell.get("scene", 0))) + if reward is None: + reward = row.get("reward") + try: + reward = float(reward) + except (TypeError, ValueError): + continue + model.train_on(cell.get("bucket", 0), cell["desc"].float(), reward) + + def train_from_rows(refinement_key, rows, n_buckets=None): """Rebuild the per-bucket heads from probe rows. -> {bucket: samples}, or None. @@ -159,21 +178,13 @@ def train_from_rows(refinement_key, rows, n_buckets=None): return None model = BucketedValue(n_buckets) - for row in usable: - per_scene = row.get("scene_rewards") or None - for cell in row["rows"]: - # Pass 0 only: a second pass runs its own schedule from its own starting sigma, - # so its bucket 2 is not the window bucket 2 steers in. - if int(cell.get("pass", 0)) != 0: - continue - reward = (per_scene or {}).get(int(cell.get("scene", 0))) - if reward is None: - reward = row.get("reward") - try: - reward = float(reward) - except (TypeError, ValueError): - continue - model.train_on(cell.get("bucket", 0), cell["desc"].float(), reward) + # ComfyUI executes a node under torch.inference_mode(), so every tensor made in here is + # an inference tensor and autograd refuses it ("Inference tensors do not track version + # counter"). Same guard the output value function's trainer carries. It has to wrap the + # whole loop, not one call: `.float()` on a stored descriptor makes a new tensor too, and + # that is the one training receives. + with torch.inference_mode(False), torch.enable_grad(): + _train_all(model, usable) try: model.save(state_path(refinement_key))