From b6319c51c49271d287ef424b9eeb5f1be3dc3ad7 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 3 Oct 2026 16:34:03 +0300 Subject: [PATCH] fix: forget fires only once the first run is queued, one owner per Generate; velocity bias banks likes only, mixed sizes import, says when no structure step, refuses windows/second-order/negative-only calls Co-Authored-By: Claude Sonnet 5.5 --- app/legacy/store.js | 17 ++++--- modules/sampling/velocity_bias/__init__.py | 46 +++++++++++++++---- .../velocity_bias/tests/test_velocity_bias.py | 46 +++++++++++++++++++ modules/system/taste/__init__.py | 8 ++-- modules/system/taste/store.py | 37 ++++++++++----- modules/system/taste/tests/test_taste.py | 30 ++++++++++++ 6 files changed, 153 insertions(+), 31 deletions(-) diff --git a/app/legacy/store.js b/app/legacy/store.js index e574db69..c55cc607 100644 --- a/app/legacy/store.js +++ b/app/legacy/store.js @@ -4019,10 +4019,6 @@ return false; } _genRunActive = true; - if (_tasteGenPending) { // the first run of a Generate: what the last one left unrated is forgotten - _tasteGenPending = false; - await _newTasteGeneration(); - } // A run started through the normal click path owns its own clock // lifecycle (the _clockedX wrapper's `finally`, spanning the whole // batch) — any stale flag left over from a PRIOR adopted run must not @@ -4074,6 +4070,13 @@ // clobber it. Setting it last means an unwired reference is always // what the person sees, never lost to a race with an unrelated note. if (unwiredNotice) { state.notice = unwiredNotice; notify(); } + if (queued && _tasteGenPending) { + // The first run of a Generate is really queued: what the last Generate left unrated is + // forgotten now (a refused or failed click above never gets here). Sampling takes seconds; + // this takes milliseconds, so it lands before this run's own capture. + _tasteGenPending = false; + await _newTasteGeneration(); + } if (!queued) { // Refused before it ever reached run.start() (an incomplete pipeline, // or the queue itself saying no) -- the bridge's own "say" hook @@ -4338,14 +4341,16 @@ } // Runs `fn` as the owner of one Generate -- unless one is already in flight, in which case this click // is refused later anyway and must not make the running Generate forget its unrated clips. + let _generateOwned = false; // an entry point holds a Generate (set before its first await) async function _ownsGenerate(fn) { - const owner = !_genRunActive; + const owner = !_genRunActive && !_generateOwned; if (owner) { + _generateOwned = true; _tasteGenPending = true; // The settings of this Generate, taken once: every run of it is queued with these. _frozenSettings = window.PipelineState && window.PipelineState.frozenInputs ? window.PipelineState.frozenInputs() : null; } - try { return await fn(); } finally { if (owner) { _tasteGenPending = false; _frozenSettings = null; } } + try { return await fn(); } finally { if (owner) { _tasteGenPending = false; _frozenSettings = null; _generateOwned = false; } } } async function _clockedGenerate(onlyScene) { _genClockStart(onlyScene ? "scene" : "all"); diff --git a/modules/sampling/velocity_bias/__init__.py b/modules/sampling/velocity_bias/__init__.py index 3b88d967..fa7eeca4 100644 --- a/modules/sampling/velocity_bias/__init__.py +++ b/modules/sampling/velocity_bias/__init__.py @@ -18,6 +18,7 @@ import torch.nn.functional as F from comfy.patcher_extension import WrappersMP from ..._core import dit_hooks, log, registry, streams +from ..score_slider import cond_row ID = "velocity_bias" TITLE = "Velocity bias" @@ -123,46 +124,71 @@ def install(patcher, values, key): return None strength = max(0.0, min(3.0, float(values.get("strength", 0.15)))) nearest = bool(values.get("nearest")) - live = {"rows": []} + live = {"rows": [], "hit": False, "seen": {}, "multi": False} def fresh(): + live["hit"], live["multi"] = False, False + live["seen"].clear() live["rows"] = taste.rows(KIND) liked = sum(1 for r in live["rows"] if r["reward"] > 0) log.once(f"{ID}:state", log.INFO, "FunPack Velocity bias", f"key {taste.key!r}: {liked} liked clip(s) banked" if liked else f"key {taste.key!r}: learning, nothing liked banked yet") - captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS, fresh=fresh) + # Only liked clips are banked (nothing here learns from a disliked one); clips of different sizes + # are meant to coexist, so a key holding two resolutions still exports and imports. + captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS, fresh=fresh, only="liked", mixed=True) def apply_model(executor, x, t, *args, **kwargs): named = streams.model_args(args, kwargs) to = named.get("transformer_options") - if dit_hooks.probing(to) or not structure_step(to): + if dit_hooks.probing(to): return executor(x, t, *args, **kwargs) + if (to or {}).get("context_window") is not None: + _say("off this run: it cannot follow clips cut into context windows (each window is a different size)") + return executor(x, t, *args, **kwargs) + where = dit_hooks.current_step(to) + if where is not None and where[0] >= where[1] - 1 and not live["hit"] and not live["multi"]: + _say("did nothing this run: no step of this schedule is near 90% of the starting noise, " + "so there was no structure step to bias or learn from") + if not structure_step(to): + return executor(x, t, *args, **kwargs) + kinds = tuple(int(v) for v in ((to or {}).get("cond_or_uncond") or ())) c = named.get("c_crossattn") + row = cond_row(c, to) if torch.is_tensor(c) and c.dim() == 3 else 0 + if row is None: # only the negative prompt in this call + return executor(x, t, *args, **kwargs) + if live["seen"].get(kinds) == where[0] or live["multi"]: + if not live["multi"]: + live["multi"] = True + _say("off for the rest of this run: the model is called more than once per step " + "(a second-order sampler); use euler-style sampling") + return executor(x, t, *args, **kwargs) + live["seen"][kinds] = where[0] sig = None if torch.is_tensor(c) and c.dim() == 3: words = registry.current().ask("text_rows", named, int(c.shape[1])) - row = c[0] if words is None else c[0][words.to(c.device)] - sig = row.float().mean(0).detach().cpu() + r = c[row] if words is None else c[row][words.to(c.device)] + sig = r.float().mean(0).detach().cpu() split_in = streams.video_of(x, named) if split_in is None: _say("off this run: could not find the picture in this model's latent") return executor(x, t, *args, **kwargs) + live["hit"] = True video_in, rebuild = split_in sigma = float(t.max()) ratio = sigma / float(to["sample_sigmas"][0]) - direction, how = reference(live["rows"], video_in[0].shape, sig, nearest) if strength > 0 else (None, "") + direction, how = reference(live["rows"], video_in[row].shape, sig, nearest) if strength > 0 else (None, "") run_x = x if direction is not None: run_x = rebuild(rotate(video_in, direction.unsqueeze(0), strength, ratio)) - log.once(f"{ID}:applied", log.INFO, "FunPack Velocity bias", f"applied: {how}") - elif strength > 0 and live["rows"]: - _say("not applied this run: " + reference(live["rows"], video_in[0].shape, sig, nearest)[1]) + log.once(f"{ID}:applied:{how}", log.INFO, "FunPack Velocity bias", f"applied: {how}") + elif strength > 0: + _say("not applied this run: " + reference(live["rows"], video_in[row].shape, sig, nearest)[1]) out = executor(run_x, t, *args, **kwargs) split_out, split_run = streams.video_of(out, named), streams.video_of(run_x, named) if split_out is not None and split_run is not None: - d = ((split_run[0].float() - split_out[0].float()) / max(sigma, 1e-6))[0] # where the model moves the picture + d = ((split_run[0].float() - split_out[0].float()) / max(sigma, 1e-6))[row] # where the model moves the picture captured["v"] = d.detach().half().cpu() if sig is not None: captured["sig"] = sig diff --git a/modules/sampling/velocity_bias/tests/test_velocity_bias.py b/modules/sampling/velocity_bias/tests/test_velocity_bias.py index bfe95a94..12ae9742 100644 --- a/modules/sampling/velocity_bias/tests/test_velocity_bias.py +++ b/modules/sampling/velocity_bias/tests/test_velocity_bias.py @@ -111,3 +111,49 @@ def test_strength_zero_only_learns(tiny_h3): seen = [] wrap(lambda x, *a, **k: seen.append(x) or x0, x0, sched[1:2], None, None, None, _to(1, sched), latent_shapes=shapes) assert torch.equal(seen[0], x0) + + +def _packed(shapes_seed=0): + from conftest import packed_av + torch.manual_seed(shapes_seed) + return packed_av(torch.randn(1, 4, 3, 8, 8), torch.randn(1, 8, 5)) + + +def test_a_schedule_with_no_structure_step_says_so_on_its_last_step(tiny_h3): + from core import log + wrap, _ = _load(tiny_h3, strength=1.0) + x0, shapes = _packed() + sched = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + log.new_run() + before = len(log.history()) + for i in range(4): + wrap(lambda x, *a, **k: x0, x0, sched[i:i + 1], None, None, None, _to(i, sched), latent_shapes=shapes) + assert any("no step of this schedule is near 90%" in e["message"] for e in log.history()[before:]) + + +def test_context_windows_and_second_order_calls_are_refused_in_words(tiny_h3): + from core import log + x0, shapes = _packed() + sched = torch.tensor([1.0, 0.909375, 0.725, 0.421875, 0.0]) + wrap, _ = _load(tiny_h3, strength=1.0) + log.new_run() + wrap(lambda x, *a, **k: x0, x0, sched[1:2], None, None, None, {**_to(1, sched), "context_window": object()}, + latent_shapes=shapes) + assert any("context windows" in e["message"] for e in log.history()) + wrap2, outer = _load(tiny_h3, strength=1.0) + log.new_run() + for o in outer[:1]: + o(lambda: None) # a run starts: the state is fresh + for _ in range(2): # the same structure step twice: a corrector call + wrap2(lambda x, *a, **k: x0, x0, sched[1:2], None, None, None, _to(1, sched), latent_shapes=shapes) + assert any("more than once per step" in e["message"] for e in log.history()) + + +def test_a_call_that_carries_only_the_negative_prompt_neither_biases_nor_learns(tiny_h3): + wrap, _ = _load(tiny_h3, strength=1.0) + x0, shapes = _packed() + sched = torch.tensor([1.0, 0.909375, 0.725, 0.421875, 0.0]) + seen = [] + to = {**_to(1, sched), "cond_or_uncond": [1]} + wrap(lambda x, *a, **k: seen.append(x) or x0, x0, sched[1:2], None, torch.ones(1, 4, 8), None, to, latent_shapes=shapes) + assert torch.equal(seen[0], x0) diff --git a/modules/system/taste/__init__.py b/modules/system/taste/__init__.py index 01512a82..5aba7eca 100644 --- a/modules/system/taste/__init__.py +++ b/modules/system/taste/__init__.py @@ -38,8 +38,8 @@ class Handle: def __init__(self, key): self.key = key - def capture(self, kind, rows, keep=store.MAX_ROWS): - store.capture(self.key, kind, rows, keep=keep) + def capture(self, kind, rows, keep=store.MAX_ROWS, only=None, mixed=False): + store.capture(self.key, kind, rows, keep=keep, only=only, mixed=mixed) def rows(self, kind, blind_to=None): """Every rated row of `kind`, oldest first: {"reward", "rows", "prompt_id"}. @@ -47,7 +47,7 @@ class Handle: neutral -- for a learner that cannot tell whether that half was the problem.""" return store.blind(store.load(self.key, kind)["rows"], blind_to) - def collect(self, patcher, key, kind, keep=store.MAX_ROWS, fresh=None): + def collect(self, patcher, key, kind, keep=store.MAX_ROWS, fresh=None, only=None, mixed=False): """A dict to fill during sampling; kept as `kind`'s capture when the sampling call finishes. An interrupted or failed run keeps nothing, so a half-finished clip can never be what a rating teaches. @@ -64,7 +64,7 @@ class Handle: if fresh is not None: fresh() out = executor(*args, **kwargs) - self.capture(kind, rows, keep=keep) + self.capture(kind, rows, keep=keep, only=only, mixed=mixed) return out patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, outer) diff --git a/modules/system/taste/store.py b/modules/system/taste/store.py index d68b9f5b..383255cd 100644 --- a/modules/system/taste/store.py +++ b/modules/system/taste/store.py @@ -134,12 +134,20 @@ def _pending_path(key, kind, prompt_id=None): return _dir(key) / (f"{kind}.pending.pt" if prompt_id is None else f"{kind}.{prompt_id}.pending.pt") -def capture(key, kind, rows, prompt_id=None, keep=MAX_ROWS): +def _prune_aside(key, kind): + aside = sorted(_dir(key).glob(f"{kind}.*.pending.pt"), key=lambda p: p.stat().st_mtime_ns) + for old in aside[:-MAX_PENDING]: + old.unlink(missing_ok=True) + + +def capture(key, kind, rows, prompt_id=None, keep=MAX_ROWS, only=None, mixed=False): """This run's capture for `kind`, waiting for a rating until the next Generate starts. The previous run's capture of the same Generate is kept aside, not overwritten. Kept in the dtype given: a banked latent stays half precision. `keep` caps how many rated rows this kind - holds, oldest dropped first. + holds, oldest dropped first. `only="liked"`: a kind that learns from liked clips alone does not + bank a disliked one (it could only push the liked ones out under `keep`). `mixed`: its rows are + legitimately of different shapes (a clip bank is matched by size), so import does not demand one. """ if not rows: return @@ -152,10 +160,8 @@ def capture(key, kind, rows, prompt_id=None, keep=MAX_ROWS): if earlier is not None and earlier != prompt_id: if _PID.match(str(earlier)): os.replace(latest, _pending_path(key, kind, earlier)) - _write(latest, {"prompt_id": prompt_id, "rows": clean, "keep": int(keep)}) - aside = sorted(_dir(key).glob(f"{kind}.*.pending.pt"), key=lambda p: p.stat().st_mtime_ns) - for old in aside[:-MAX_PENDING]: - old.unlink(missing_ok=True) + _write(latest, {"prompt_id": prompt_id, "rows": clean, "keep": int(keep), "only": only, "mixed": bool(mixed)}) + _prune_aside(key, kind) state = _read_latest() runs = state.get("runs", {}) kinds = runs.setdefault(prompt_id, {}).setdefault(key, []) @@ -195,7 +201,7 @@ def blind(rows, to): return [dict(r, reward=0.0) if r.get("axis") == to else r for r in rows] -def _restore_pending(key, kind, prompt_id, rows, keep): +def _restore_pending(key, kind, prompt_id, rows, keep, only=None, mixed=False): """A cleared rating puts the capture back to waiting (until the next Generate), so the clip can be rated again. The latest slot if it is free or already this run's, else aside under the prompt id.""" latest = _pending_path(key, kind) @@ -204,7 +210,8 @@ def _restore_pending(key, kind, prompt_id, rows, keep): if not _PID.match(str(prompt_id)): return latest = _pending_path(key, kind, prompt_id) - _write(latest, {"prompt_id": prompt_id, "rows": rows, "keep": int(keep)}) + _write(latest, {"prompt_id": prompt_id, "rows": rows, "keep": int(keep), "only": only, "mixed": bool(mixed)}) + _prune_aside(key, kind) def rate(prompt_id, rating, axis=None): @@ -237,7 +244,8 @@ def rate(prompt_id, rating, axis=None): if rating is None: data["rows"] = [r for r in data["rows"] if r.get("prompt_id") != prompt_id] if path.stem in entry.get(key, []): - _restore_pending(key, path.stem, prompt_id, hit[-1]["rows"], hit[-1].get("keep", MAX_ROWS)) + _restore_pending(key, path.stem, prompt_id, hit[-1]["rows"], hit[-1].get("keep", MAX_ROWS), + hit[-1].get("only"), hit[-1].get("mixed", False)) else: for r in hit: r["reward"] = REWARD[rating] @@ -269,7 +277,12 @@ def rate(prompt_id, rating, axis=None): continue data = load(key, kind) keep = int(pending.get("keep", MAX_ROWS)) - row = {"prompt_id": prompt_id, "reward": REWARD[rating], "rows": pending["rows"], "keep": keep} + if pending.get("only") == "liked" and rating != "liked": + pending_path.unlink(missing_ok=True) # nothing here learns from it: not banked + out.setdefault("skipped", []).append(kind) + continue + row = {"prompt_id": prompt_id, "reward": REWARD[rating], "rows": pending["rows"], "keep": keep, + "only": pending.get("only"), "mixed": bool(pending.get("mixed"))} if axis: row["axis"] = axis data["rows"].append(row) @@ -277,7 +290,7 @@ def rate(prompt_id, rating, axis=None): _write(_dir(key) / f"{kind}.pt", data) pending_path.unlink(missing_ok=True) out["recorded"].append(kind) - if not out["recorded"]: + if not out["recorded"] and not out.get("skipped"): out["why"] = ("this clip's capture was not kept (more runs waited than are kept, or it was " "cleared); nothing was learned") return out @@ -384,6 +397,8 @@ def _sound_rows(data): return False # What is stacked later must agree: one name, one shape (a key trained across two model # families would otherwise abort every sampling that reads it). + if row.get("mixed") is True: + continue # a clip bank: rows of different sizes are the point for name, v in captured.items(): if shapes.setdefault(name, tuple(v.shape)) != tuple(v.shape): return False diff --git a/modules/system/taste/tests/test_taste.py b/modules/system/taste/tests/test_taste.py index 276ed061..ea6ac7eb 100644 --- a/modules/system/taste/tests/test_taste.py +++ b/modules/system/taste/tests/test_taste.py @@ -210,3 +210,33 @@ def test_a_run_waiting_under_two_keys_teaches_both(): store.capture("wolf", "reins", {25: _v(2)}, prompt_id="p1") assert store.rate("p1", "liked")["recorded"] == ["reins", "reins"] assert store.counts("fox", "reins") == (1, 0) and store.counts("wolf", "reins") == (1, 0) + + +def test_a_liked_only_kind_banks_likes_and_never_lets_dislikes_push_them_out(): + for i in range(3): + store.capture("fox", "vel", {"v": _v(i)}, prompt_id=f"l{i}", keep=4, only="liked") + assert store.rate(f"l{i}", "liked")["recorded"] == ["vel"] + for i in range(10): + store.capture("fox", "vel", {"v": _v(9)}, prompt_id=f"d{i}", keep=4, only="liked") + out = store.rate(f"d{i}", "disliked") + assert out["recorded"] == [] and out["skipped"] == ["vel"] and out["why"] is None + assert store.counts("fox", "vel") == (3, 0) + + +def test_a_kind_of_mixed_sizes_exports_and_imports(tmp_path): + store.capture("fox", "vel", {"v": torch.ones(2, 2)}, prompt_id="a", mixed=True) + store.rate("a", "liked") + store.capture("fox", "vel", {"v": torch.ones(3, 3)}, prompt_id="b", mixed=True) + store.rate("b", "liked") + out = tmp_path / "k.zip" + store.export_key("fox", out) + assert store.import_key("wolf", out) == 1 + assert store.counts("wolf", "vel") == (2, 0) + # a kind that did not say so still must hold one shape per name + store.capture("bear", "reins", {"v": torch.ones(2, 2)}, prompt_id="c") + store.rate("c", "liked") + store.capture("bear", "reins", {"v": torch.ones(3, 3)}, prompt_id="d") + store.rate("d", "liked") + store.export_key("bear", tmp_path / "b.zip") + with pytest.raises(ValueError): + store.import_key("cub", tmp_path / "b.zip")